amoghrajesh commented on code in PR #76033:
URL: https://github.com/apache/airflow/pull/76033#discussion_r4440192190
##########
providers/amazon/tests/unit/amazon/aws/operators/test_sagemaker_training.py:
##########
@@ -116,6 +116,52 @@ def test_execute_without_check_if_job_exists(self, _,
mock_training, mock_desc):
a = []
a.sort()
+ @mock.patch.object(SageMakerHook, "stop_training_job")
+ def test_on_kill_does_nothing_by_default(self, mock_stop):
+ self.sagemaker._training_job_name = "job_name"
+
+ self.sagemaker.on_kill()
+
+ mock_stop.assert_not_called()
+
+ @mock.patch.object(SageMakerHook, "stop_training_job")
+ def test_on_kill_does_nothing_before_the_job_is_submitted(self, mock_stop):
+ self.sagemaker.stop_training_job_on_kill = True
+
+ self.sagemaker.on_kill()
+
+ mock_stop.assert_not_called()
+
+ @mock.patch.object(SageMakerHook, "stop_training_job")
+ def test_on_kill_stops_the_job_when_enabled(self, mock_stop):
+ self.sagemaker.stop_training_job_on_kill = True
+ self.sagemaker._training_job_name = "job_name"
+
+ self.sagemaker.on_kill()
+
+ mock_stop.assert_called_once_with("job_name")
+
+ @mock.patch.object(SageMakerHook, "stop_training_job")
+ @mock.patch.object(SageMakerHook, "describe_training_job")
+ @mock.patch.object(SageMakerHook, "create_training_job")
+ @mock.patch.object(sagemaker, "serialize", return_value="")
+ def test_on_kill_stops_the_renamed_job(self, _, mock_training, mock_desc,
mock_stop):
Review Comment:
Replaced `test_on_kill_stops_the_job_when_enabled` with
`test_on_kill_stops_a_job_killed_during_submit`, which uses
`wait_for_completion=True` and a `create_training_job` side effect that calls
`on_kill()` mid-submit, and builds the operator with
`stop_training_job_on_kill=True` through the constructor. Confirmed it fails
when the `_training_job_name` assignment moves below `create_training_job`,
while the other `on_kill` tests still pass.
##########
providers/amazon/tests/unit/amazon/aws/hooks/test_sagemaker.py:
##########
@@ -729,6 +729,40 @@ def test_start_pipeline_returns_arn(self, mock_conn):
# Value contains the value associated with the key in Name
assert transformed_param["Value"] ==
params_dict[transformed_param["Name"]]
+ @patch("airflow.providers.amazon.aws.hooks.sagemaker.SageMakerHook.conn",
new_callable=mock.PropertyMock)
+ def test_stop_training_job(self, mock_conn):
+ hook = SageMakerHook(aws_conn_id="aws_default")
+
+ hook.stop_training_job(training_job_name="job_name")
+
+
mock_conn().stop_training_job.assert_called_once_with(TrainingJobName="job_name")
+
+ @patch("airflow.providers.amazon.aws.hooks.sagemaker.SageMakerHook.conn",
new_callable=mock.PropertyMock)
+ def test_stop_training_job_swallows_validation_exception(self, mock_conn):
+ mock_conn().stop_training_job.side_effect = ClientError(
+ error_response={"Error": {"Code": "ValidationException",
"Message": "already Completed"}},
+ operation_name="StopTrainingJob",
+ )
+
+ hook = SageMakerHook(aws_conn_id="aws_default")
+
+ hook.stop_training_job(training_job_name="job_name")
Review Comment:
Sure
--
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]
