Fix unexpected bug in exiting hook context manager (#18014)

apache · Sep 3, 2021 · 9c64419 · 9c64419
1 parent a88115e
commit 9c64419
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 1 deletion.
diff --git a/airflow/providers/microsoft/psrp/hooks/psrp.py b/airflow/providers/microsoft/psrp/hooks/psrp.py
@@ -57,7 +57,7 @@ def __enter__(self):
 
     def __exit__(self, exc_type, exc_value, traceback):
         try:
-            self._client.__exit__()
+            self._client.__exit__(exc_type, exc_value, traceback)
         finally:
             self._client = None
 

diff --git a/tests/providers/microsoft/psrp/hooks/test_psrp.py b/tests/providers/microsoft/psrp/hooks/test_psrp.py
@@ -64,6 +64,8 @@ def end_invoke():
 
             hook.invoke_powershell("foo")
 
+        assert ws_man().__exit__.mock_calls == [call(None, None, None)]
+
         assert call('%s', '<output>') in log_info.mock_calls
         assert call('Information: %s', '<message>') in log_info.mock_calls
         assert call('Invocation state: %s', 'Completed') in log_info.mock_calls