josh-fell commented on code in PR #36303:
URL: https://github.com/apache/airflow/pull/36303#discussion_r1431933045
##########
tests/providers/ssh/operators/test_ssh.py:
##########
@@ -203,6 +203,37 @@ def test_ssh_client_managed_correctly(self):
self.hook.get_conn.assert_called_once()
self.hook.get_conn.return_value.__exit__.assert_called_once()
+ @pytest.mark.parametrize(
+ "extra_kwargs, actual_exit_code, expected_exc",
+ [
+ ({}, 0, None),
+ ({}, 100, AirflowException),
+ ({"skip_on_exit_code": 100}, 100, AirflowSkipException),
+ ({"skip_on_exit_code": 100}, 101, AirflowException),
+ ({"skip_on_exit_code": None}, 100, AirflowException),
+ ({"skip_on_exit_code": [100]}, 100, AirflowSkipException),
+ ({"skip_on_exit_code": (100, 101)}, 100, AirflowSkipException),
+ ({"skip_on_exit_code": 100}, 101, AirflowException),
+ ({"skip_on_exit_code": [100, 102]}, 101, AirflowException),
+ ],
+ )
+ def test_skip(self, extra_kwargs, actual_exit_code, expected_exc):
+ command = "not_a_real_command"
+ self.exec_ssh_client_command.return_value = (actual_exit_code, b"",
b"")
+
+ operator = SSHOperator(
+ task_id="test",
+ ssh_hook=self.hook,
+ command=command,
+ **extra_kwargs,
+ )
+
+ if expected_exc is None:
+ operator.execute({})
+ else:
+ with pytest.raises(expected_exc):
+ operator.execute({})
Review Comment:
```suggestion
with pytest.raises(expected_exc):
operator.execute({})
```
It's a nit really, but you don't need the branch here.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
