def test_sensor(self, mock_describe_job, hook_init, mock_client): hook_init.return_value = None mock_describe_job.side_effect = [ DESCRIBE_TUNING_INPROGRESS_RESPONSE, DESCRIBE_TUNING_STOPPING_RESPONSE, DESCRIBE_TUNING_COMPELETED_RESPONSE ] sensor = SageMakerTuningSensor(task_id='test_task', poke_interval=2, aws_conn_id='aws_test', job_name='test_job_name') 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_calls_until_a_terminal_state(self, mock_describe_job, hook_init, mock_client): hook_init.return_value = None mock_describe_job.side_effect = [ DESCRIBE_TUNING_INPROGRESS_RETURN, DESCRIBE_TUNING_STOPPING_RETURN, DESCRIBE_TUNING_STOPPED_RETURN, DESCRIBE_TUNING_COMPELETED_RETURN ] sensor = SageMakerTuningSensor( 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')
def test_sensor_with_failure(self, mock_describe_job, mock_client): mock_describe_job.side_effect = [DESCRIBE_TUNING_FAILED_RESPONSE] sensor = SageMakerTuningSensor(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_calls_until_a_terminal_state(self, mock_describe_job, hook_init, mock_client): hook_init.return_value = None mock_describe_job.side_effect = [ DESCRIBE_TUNING_INPROGRESS_RETURN, DESCRIBE_TUNING_STOPPING_RETURN, DESCRIBE_TUNING_STOPPED_RETURN, DESCRIBE_TUNING_COMPELETED_RETURN ] sensor = SageMakerTuningSensor(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')