diff --git a/providers/microsoft/psrp/src/airflow/providers/microsoft/psrp/operators/psrp.py b/providers/microsoft/psrp/src/airflow/providers/microsoft/psrp/operators/psrp.py index 4b64dc6dac578..673c99f4a37b8 100644 --- a/providers/microsoft/psrp/src/airflow/providers/microsoft/psrp/operators/psrp.py +++ b/providers/microsoft/psrp/src/airflow/providers/microsoft/psrp/operators/psrp.py @@ -106,12 +106,11 @@ def __init__( psrp_session_init: Command | None = None, **kwargs, ) -> None: - args = {command, powershell, cmdlet} - if not exactly_one(*args): + if not exactly_one(command is not None, powershell is not None, cmdlet is not None): raise ValueError("Must provide exactly one of 'command', 'powershell', or 'cmdlet'") - if arguments and not (powershell or cmdlet): + if arguments is not None and powershell is None and cmdlet is None: raise ValueError("Arguments only allowed with 'powershell' or 'cmdlet'") - if parameters and not (powershell or cmdlet): + if parameters is not None and powershell is None and cmdlet is None: raise ValueError("Parameters only allowed with 'powershell' or 'cmdlet'") if cmdlet: kwargs.setdefault("task_id", cmdlet) @@ -140,10 +139,10 @@ def execute(self, context: Context) -> list[Any] | None: ): if self.psrp_session_init is not None: ps.add_command(self.psrp_session_init) - if self.command: + if self.command is not None: ps.add_script(f"cmd.exe /c @'\n{self.command}\n'@") else: - if self.cmdlet: + if self.cmdlet is not None: ps.add_cmdlet(self.cmdlet) else: ps.add_script(self.powershell) diff --git a/providers/microsoft/psrp/tests/unit/microsoft/psrp/operators/test_psrp.py b/providers/microsoft/psrp/tests/unit/microsoft/psrp/operators/test_psrp.py index 8d2fb922d6ed2..92039ec6f5612 100644 --- a/providers/microsoft/psrp/tests/unit/microsoft/psrp/operators/test_psrp.py +++ b/providers/microsoft/psrp/tests/unit/microsoft/psrp/operators/test_psrp.py @@ -39,15 +39,51 @@ class ExecuteParameter(NamedTuple): class TestPsrpOperator: - def test_no_command_or_powershell(self): - exception_msg = "Must provide exactly one of 'command', 'powershell', or 'cmdlet'" - with pytest.raises(ValueError, match=exception_msg): - PsrpOperator(task_id="test_task_id", psrp_conn_id=CONNECTION_ID) + @pytest.mark.parametrize( + ("kwargs", "match"), + [ + pytest.param({}, "exactly one", id="no-option"), + pytest.param({"command": "", "powershell": "Get-Bar"}, "exactly one", id="two-options-one-empty"), + pytest.param({"command": "x", "powershell": "x"}, "exactly one", id="two-equal-options"), + pytest.param( + {"command": "hostname", "arguments": ["x"]}, "Arguments only allowed", id="arguments" + ), + pytest.param( + {"command": "hostname", "parameters": {"k": "v"}}, "Parameters only allowed", id="parameters" + ), + ], + ) + def test_invalid_option_combination(self, kwargs, match): + with pytest.raises(ValueError, match=match): + PsrpOperator(task_id="test_task_id", psrp_conn_id=CONNECTION_ID, **kwargs) + + @pytest.mark.parametrize( + "kwargs", + [ + pytest.param({"command": ""}, id="empty-command"), + pytest.param({"powershell": "", "arguments": ["a"]}, id="arguments"), + pytest.param({"powershell": "", "parameters": {"k": "v"}}, id="parameters"), + ], + ) + def test_empty_option_counts_as_provided(self, kwargs): + PsrpOperator(task_id="test_task_id", psrp_conn_id=CONNECTION_ID, **kwargs) def test_cmdlet_task_id_default(self): operator = PsrpOperator(cmdlet="Invoke-Foo", psrp_conn_id=CONNECTION_ID) assert operator.task_id == "Invoke-Foo" + @patch(f"{PsrpOperator.__module__}.PsrpHook") + def test_command_rendering_to_empty_dispatches_as_command(self, hook_impl): + op = PsrpOperator(task_id="test", psrp_conn_id=CONNECTION_ID, command="{{ '' }}") + op.render_template_fields({}) + assert op.command == "" + ps = Mock(spec=PowerShell, output=[], had_errors=False, runspace_pool=Mock(host=Mock(rc=0))) + hook_impl.configure_mock( + **{"return_value.__enter__.return_value.invoke.return_value.__enter__.return_value": ps} + ) + op.execute(None) + ps.add_script.assert_called_once_with("cmd.exe /c @'\n\n'@") + @pytest.mark.parametrize("do_xcom_push", [True, False]) @pytest.mark.parametrize( ("had_errors", "rc"), [(False, 0), (False, None), (True, None), (False, 1), (True, 1)]