def test_describe_tuning_job(self, mock_client):
mock_session = mock.Mock()
attrs = {'describe_hyper_parameter_tuning_job.return_value': 'InProgress'}
mock_session.configure_mock(**attrs)
mock_client.return_value = mock_session
hook = SageMakerHook(aws_conn_id='sagemaker_test_conn_id')
response = hook.describe_tuning_job(job_name)
mock_session.describe_hyper_parameter_tuning_job.assert_called_once_with(
HyperParameterTuningJobName=job_name
)
assert response == 'InProgress'
def get_sagemaker_response(self):
sagemaker = SageMakerHook(aws_conn_id=self.aws_conn_id)
self.log.info('Poking Sagemaker Tuning Job %s', self.job_name)
return sagemaker.describe_tuning_job(self.job_name)
