amoghrajesh commented on code in PR #76036:
URL: https://github.com/apache/airflow/pull/76036#discussion_r4440351951
##########
providers/amazon/tests/unit/amazon/aws/operators/test_sagemaker_training.py:
##########
@@ -79,42 +79,127 @@ def setup_method(self):
check_interval=5,
)
+ @mock.patch.object(SageMakerHook, "check_training_status_with_log")
@mock.patch.object(SageMakerHook, "describe_training_job")
@mock.patch.object(SageMakerHook, "create_training_job")
@mock.patch.object(sagemaker, "serialize", return_value="")
- def test_integer_fields(self, _, mock_training, mock_desc):
+ def test_integer_fields(self, _, mock_training, mock_desc, __):
mock_desc.side_effect = [ClientError({"Error": {"Code":
"ValidationException"}}, "op"), None]
mock_training.return_value = {
"TrainingJobArn": "test_arn",
"ResponseMetadata": {"HTTPStatusCode": 200},
}
- self.sagemaker.execute(None)
+ self.sagemaker.execute({})
assert self.sagemaker.integer_fields == EXPECTED_INTEGER_FIELDS
for key1, key2 in EXPECTED_INTEGER_FIELDS:
assert self.sagemaker.config[key1][key2] ==
int(self.sagemaker.config[key1][key2])
+ @mock.patch.object(SageMakerHook, "check_training_status_with_log")
@mock.patch.object(SageMakerHook, "describe_training_job")
@mock.patch.object(SageMakerHook, "create_training_job")
@mock.patch.object(sagemaker, "serialize", return_value="")
- def test_execute_without_check_if_job_exists(self, _, mock_training,
mock_desc):
+ def test_execute_without_check_if_job_exists(self, _, mock_training,
mock_desc, __):
mock_training.return_value = {
"TrainingJobArn": "test_arn",
"ResponseMetadata": {"HTTPStatusCode": 200},
}
self.sagemaker.check_if_job_exists = False
self.sagemaker._check_if_job_exists = mock.MagicMock()
- self.sagemaker.execute(None)
+ self.sagemaker.execute({})
self.sagemaker._check_if_job_exists.assert_not_called()
mock_training.assert_called_once_with(
CREATE_TRAINING_PARAMS,
wait_for_completion=False,
- print_log=True,
+ print_log=False,
check_interval=5,
max_ingestion_time=None,
)
- a = []
- a.sort()
+ @staticmethod
+ def _task_state_store(stored_job_name=None):
+ store = mock.MagicMock()
Review Comment:
Both applied
--
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]
