def test_sensor(self, mock_describe_job, hook_init, mock_client): hook_init.return_value = None mock_describe_job.side_effect = [ DESCRIBE_TRANSFORM_INPROGRESS_RESPONSE, DESCRIBE_TRANSFORM_STOPPING_RESPONSE, DESCRIBE_TRANSFORM_COMPLETED_RESPONSE ] sensor = SageMakerTransformSensor(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 completed) 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_failure(self, mock_describe_job, mock_client): mock_describe_job.side_effect = [DESCRIBE_TRANSFORM_FAILED_RESPONSE] sensor = SageMakerTransformSensor( 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')