PowerShell Remoting fail on non-zero exitcode

airflow/providers/microsoft/psrp/hooks/psrp.py

@@ -22,6 +22,7 @@
 from typing import Any, Callable, Dict, Iterator, Optional
 from weakref import WeakKeyDictionary
 
+from pypsrp.host import PSHost
 from pypsrp.messages import MessageType
 from pypsrp.powershell import PowerShell, PSInvocationState, RunspacePool
 from pypsrp.wsman import WSMan
@@ -64,6 +65,9 @@ class PsrpHook(BaseHook):
     :param exchange_keys:
         If true (default), automatically initiate a session key exchange when the
         hook is used as a context manager.
+    :param host:
+        Optional PowerShell host instance. If this is not set, the default
+        implementation will be used.
 
     You can provide an alternative `configuration_name` using either `runspace_options`
     or by setting this key as the extra fields of your connection.
@@ -82,6 +86,7 @@ def __init__(
         wsman_options: Optional[Dict[str, Any]] = None,
         on_output_callback: Optional[OutputCallback] = None,
         exchange_keys: bool = True,
+        host: Optional[PSHost] = None,
     ):
         self.conn_id = psrp_conn_id
         self._logging_level = logging_level
@@ -90,6 +95,7 @@ def __init__(
         self._wsman_options = wsman_options or {}
         self._on_output_callback = on_output_callback
         self._exchange_keys = exchange_keys
+        self._host = host or PSHost(None, None, False, type(self).__name__, None, None, "1.0")
 
     def __enter__(self):
         conn = self.get_conn()
@@ -144,7 +150,7 @@ def apply_extra(d, keys):
 
         if extra:
             raise AirflowException(f"Unexpected extra configuration keys: {', '.join(sorted(extra))}")
-        pool = RunspacePool(wsman, **runspace_options)
+        pool = RunspacePool(wsman, host=self._host, **runspace_options)
         self._wsman_ref[pool] = wsman
         return pool
 

airflow/providers/microsoft/psrp/operators/psrp.py

@@ -146,6 +146,10 @@ def execute(self, context: "Context") -> Optional[List[Any]]:
         if ps.had_errors:
             raise AirflowException("Process failed")
 
+        rc = ps.runspace_pool.host.rc
+        if rc:
+            raise AirflowException(f"Process exited with non-zero status code: {rc}")
+
         if not self.do_xcom_push:
             return None
 

tests/providers/microsoft/psrp/hooks/test_psrp.py

@@ -22,6 +22,7 @@
 from unittest.mock import MagicMock, Mock, call, patch
 
 from parameterized import parameterized
+from pypsrp.host import PSHost
 from pypsrp.messages import MessageType
 from pypsrp.powershell import PSInvocationState
 from pytest import raises
@@ -180,8 +181,10 @@ def assert_log(level, *args):
         assert_log(INFO, DUMMY_STACKTRACE[1])
 
         assert call('Invocation state: %s', 'Completed') in logger.info.mock_calls
-
-        assert runspace_pool.call_args == call(ws_man.return_value, connection_name='foo')
+        args, kwargs = runspace_pool.call_args
+        assert args == (ws_man.return_value,)
+        assert kwargs["connection_name"] == "foo"
+        assert isinstance(kwargs["host"], PSHost)
 
     def test_invoke_cmdlet(self, *mocks):
         with PsrpHook(CONNECTION_ID) as hook:

tests/providers/microsoft/psrp/operators/test_psrp.py

@@ -60,13 +60,20 @@ def test_cmdlet_task_id_default(self):
                     ExecuteParameter("powershell", call.add_script("foo"), None),
                     ExecuteParameter("cmdlet", call.add_cmdlet("foo"), {"bar": "baz"}),
                 ],
-                [False, True],
+                [
+                    (False, 0),
+                    (False, None),
+                    (True, None),
+                    (False, 1),
+                    (True, 1),
+                ],
                 [False, True],
             )
         )
     )
     @patch(f"{PsrpOperator.__module__}.PsrpHook")
-    def test_execute(self, parameter, had_errors, do_xcom_push, hook_impl):
+    def test_execute(self, parameter, result, do_xcom_push, hook_impl):
+        had_errors, rc = result
         kwargs = {parameter.name: "foo"}
         if parameter.expected_parameters:
             kwargs["parameters"] = parameter.expected_parameters
@@ -78,12 +85,19 @@ def test_execute(self, parameter, had_errors, do_xcom_push, hook_impl):
             do_xcom_push=do_xcom_push,
             **kwargs,
         )
-        ps = Mock(spec=PowerShell, output=[json.dumps("<output>")], had_errors=had_errors)
+        runspace_pool = Mock()
+        runspace_pool.host.rc = rc
+        ps = Mock(
+            spec=PowerShell,
+            output=[json.dumps("<output>")],
+            had_errors=had_errors,
+            runspace_pool=runspace_pool,
+        )
         hook_impl.configure_mock(
             **{"return_value.__enter__.return_value.invoke.return_value.__enter__.return_value": ps}
         )
-        if had_errors:
-            exception_msg = "Process failed"
+        if had_errors or rc:
+            exception_msg = "Process failed" if had_errors else "Process exited with non-zero status code: 1"
             with pytest.raises(AirflowException, match=exception_msg):
                 op.execute(None)
         else: