def test_sensor(self, mock_describe_job, hook_init, mock_client):
hook_init.return_value = None
mock_describe_job.side_effect = [
DESCRIBE_TRAINING_INPROGRESS_RESPONSE,
DESCRIBE_TRAINING_STOPPING_RESPONSE,
DESCRIBE_TRAINING_COMPELETED_RESPONSE
]
sensor = SageMakerTrainingSensor(
task_id='test_task',
poke_interval=2,
aws_conn_id='aws_test',
job_name='test_job_name',
print_log=False
)
sensor.execute(None)
# make sure we called 3 times(terminated when its compeleted)
self.assertEqual(mock_describe_job.call_count, 3)
# make sure the hook was initialized with the specific params
calls = [
mock.call(aws_conn_id='aws_test'),
mock.call(aws_conn_id='aws_test'),
mock.call(aws_conn_id='aws_test')
]
hook_init.assert_has_calls(calls)
def test_sensor_with_log(self, mock_describe_job, mock_describe_job_with_log,
hook_init, mock_log_client, mock_client):
hook_init.return_value = None
mock_describe_job.return_value = DESCRIBE_TRAINING_COMPELETED_RESPONSE
mock_describe_job_with_log.side_effect = [
(LogState.WAIT_IN_PROGRESS, DESCRIBE_TRAINING_INPROGRESS_RESPONSE, 0),
(LogState.JOB_COMPLETE, DESCRIBE_TRAINING_STOPPING_RESPONSE, 0),
(LogState.COMPLETE, DESCRIBE_TRAINING_COMPELETED_RESPONSE, 0)
]
sensor = SageMakerTrainingSensor(
task_id='test_task',
poke_interval=2,
aws_conn_id='aws_test',
job_name='test_job_name',
print_log=True
)
sensor.execute(None)
self.assertEqual(mock_describe_job_with_log.call_count, 3)
self.assertEqual(mock_describe_job.call_count, 1)
calls = [
mock.call(aws_conn_id='aws_test'),
mock.call(aws_conn_id='aws_test'),
mock.call(aws_conn_id='aws_test')
]
hook_init.assert_has_calls(calls)
def test_raises_errors_failed_state(self, mock_describe_job, mock_client):
mock_describe_job.side_effect = [DESCRIBE_TRAINING_FAILED_RETURN]
sensor = SageMakerTrainingSensor(task_id='test_task',
poke_interval=2,
aws_conn_id='aws_test',
job_name='test_job_name')
self.assertRaises(AirflowException, sensor.execute, None)
mock_describe_job.assert_called_once_with('test_job_name')
def test_sensor_with_failure(self, mock_describe_job, hook_init,
mock_client):
hook_init.return_value = None
mock_describe_job.side_effect = [DESCRIBE_TRAINING_FAILED_RESPONSE]
sensor = SageMakerTrainingSensor(task_id='test_task',
poke_interval=2,
aws_conn_id='aws_test',
job_name='test_job_name',
print_log=False)
self.assertRaises(AirflowException, sensor.execute, None)
mock_describe_job.assert_called_once_with('test_job_name')
def test_sensor(self, mock_describe_job, hook_init, mock_client):
hook_init.return_value = None
mock_describe_job.side_effect = [
DESCRIBE_TRAINING_INPROGRESS_RESPONSE,
DESCRIBE_TRAINING_STOPPING_RESPONSE,
DESCRIBE_TRAINING_COMPELETED_RESPONSE
]
sensor = SageMakerTrainingSensor(
task_id='test_task',
poke_interval=2,
aws_conn_id='aws_test',
job_name='test_job_name',
print_log=False
)
sensor.execute(None)
# make sure we called 3 times(terminated when its compeleted)
self.assertEqual(mock_describe_job.call_count, 3)
# make sure the hook was initialized with the specific params
hook_init.assert_called_with(aws_conn_id='aws_test')
def test_sensor_with_log(self, mock_describe_job, mock_describe_job_with_log,
hook_init, mock_log_client, mock_client):
hook_init.return_value = None
mock_describe_job.return_value = DESCRIBE_TRAINING_COMPELETED_RESPONSE
mock_describe_job_with_log.side_effect = [
(LogState.WAIT_IN_PROGRESS, DESCRIBE_TRAINING_INPROGRESS_RESPONSE, 0),
(LogState.JOB_COMPLETE, DESCRIBE_TRAINING_STOPPING_RESPONSE, 0),
(LogState.COMPLETE, DESCRIBE_TRAINING_COMPELETED_RESPONSE, 0)
]
sensor = SageMakerTrainingSensor(
task_id='test_task',
poke_interval=2,
aws_conn_id='aws_test',
job_name='test_job_name',
print_log=True
)
sensor.execute(None)
self.assertEqual(mock_describe_job_with_log.call_count, 3)
self.assertEqual(mock_describe_job.call_count, 1)
hook_init.assert_called_with(aws_conn_id='aws_test')
def test_calls_until_a_terminal_state(self, mock_describe_job, hook_init,
mock_client):
hook_init.return_value = None
mock_describe_job.side_effect = [
DESCRIBE_TRAINING_INPROGRESS_RETURN,
DESCRIBE_TRAINING_STOPPING_RETURN,
DESCRIBE_TRAINING_STOPPED_RETURN,
DESCRIBE_TRAINING_COMPELETED_RETURN
]
sensor = SageMakerTrainingSensor(task_id='test_task',
poke_interval=2,
aws_conn_id='aws_test',
job_name='test_job_name',
region_name='us-east-1')
sensor.execute(None)
# make sure we called 4 times(terminated when its compeleted)
self.assertEqual(mock_describe_job.call_count, 4)
# make sure the hook was initialized with the specific params
hook_init.assert_called_with(aws_conn_id='aws_test',
region_name='us-east-1')