amoghrajesh commented on code in PR #76413:
URL: https://github.com/apache/airflow/pull/76413#discussion_r4508683715
##########
providers/amazon/tests/unit/amazon/aws/operators/test_sagemaker_processing.py:
##########
@@ -112,6 +112,77 @@ def setup_method(self):
task_id="test_sagemaker_operator", wait_for_completion=True,
check_interval=5, deferrable=True
)
+ @mock.patch.object(SageMakerHook, "stop_processing_job")
+ def test_on_kill_does_nothing_by_default(self, mock_stop):
+ op = SageMakerProcessingOperator(**self.processing_config_kwargs,
config=CREATE_PROCESSING_PARAMS)
+ op._processing_job_name = "job_name"
+
+ op.on_kill()
+
+ mock_stop.assert_not_called()
+
+ @mock.patch.object(SageMakerHook, "stop_processing_job")
+ def test_on_kill_does_nothing_before_the_job_is_submitted(self, mock_stop):
+ op = SageMakerProcessingOperator(
+ **self.processing_config_kwargs,
+ config=CREATE_PROCESSING_PARAMS,
+ stop_processing_job_on_kill=True,
+ )
+
+ op.on_kill()
+
+ mock_stop.assert_not_called()
+
+ @mock.patch.object(SageMakerHook, "stop_processing_job")
+ @mock.patch.object(SageMakerHook, "describe_processing_job")
+ @mock.patch.object(SageMakerHook, "count_processing_jobs_by_name",
return_value=0)
+ @mock.patch.object(SageMakerHook, "create_processing_job")
+ @mock.patch.object(sagemaker, "serialize", return_value="")
+ def test_on_kill_stops_a_job_killed_during_submit(self, _, mock_create,
__, mock_desc, mock_stop):
+ # Only holds while the name is recorded before the submit call.
+ mock_desc.side_effect = [ClientError({"Error": {"Code":
"ValidationException"}}, "op"), None]
+ op = SageMakerProcessingOperator(
+ task_id="test_sagemaker_operator",
+ config=CREATE_PROCESSING_PARAMS,
+ wait_for_completion=True,
+ check_interval=5,
+ stop_processing_job_on_kill=True,
+ )
+ mock_create.side_effect = lambda *args, **kwargs: (
+ op.on_kill(),
+ {"ProcessingJobArn": "test_arn", "ResponseMetadata":
{"HTTPStatusCode": 200}},
+ )[1]
+
+ op.execute(None)
+
+ mock_stop.assert_called_once_with("job_name")
Review Comment:
Ported as `test_on_kill_stops_the_renamed_job`. Confirmed it catches the
case: with the assignment moved above the rename, the log shows the job renamed
to `job_name-1794187155` while `on_kill` stopped `job_name`, and only this test
fails.
--
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]
