class TestEmrStepSensor(unittest.TestCase): def setUp(self): self.emr_client_mock = MagicMock() self.sensor = EmrStepSensor( task_id='test_task', poke_interval=0, job_flow_id='j-8989898989', step_id='s-VK57YR1Z9Z5N', aws_conn_id='aws_default', ) mock_emr_session = MagicMock() mock_emr_session.client.return_value = self.emr_client_mock # Mock out the emr_client creator self.boto3_session_mock = MagicMock(return_value=mock_emr_session) def test_step_completed(self): self.emr_client_mock.describe_step.side_effect = [ DESCRIBE_JOB_STEP_RUNNING_RETURN, DESCRIBE_JOB_STEP_COMPLETED_RETURN ] with patch('boto3.session.Session', self.boto3_session_mock): self.sensor.execute(None) self.assertEqual(self.emr_client_mock.describe_step.call_count, 2) calls = [ unittest.mock.call(ClusterId='j-8989898989', StepId='s-VK57YR1Z9Z5N'), unittest.mock.call(ClusterId='j-8989898989', StepId='s-VK57YR1Z9Z5N') ] self.emr_client_mock.describe_step.assert_has_calls(calls) def test_step_cancelled(self): self.emr_client_mock.describe_step.side_effect = [ DESCRIBE_JOB_STEP_RUNNING_RETURN, DESCRIBE_JOB_STEP_CANCELLED_RETURN ] with patch('boto3.session.Session', self.boto3_session_mock): self.assertRaises(AirflowException, self.sensor.execute, None) def test_step_failed(self): self.emr_client_mock.describe_step.side_effect = [ DESCRIBE_JOB_STEP_RUNNING_RETURN, DESCRIBE_JOB_STEP_FAILED_RETURN ] with patch('boto3.session.Session', self.boto3_session_mock): self.assertRaises(AirflowException, self.sensor.execute, None) def test_step_interrupted(self): self.emr_client_mock.describe_step.side_effect = [ DESCRIBE_JOB_STEP_RUNNING_RETURN, DESCRIBE_JOB_STEP_INTERRUPTED_RETURN ] with patch('boto3.session.Session', self.boto3_session_mock): self.assertRaises(AirflowException, self.sensor.execute, None)
class TestEmrStepSensor(unittest.TestCase): def setUp(self): configuration.load_test_config() self.emr_client_mock = MagicMock() self.sensor = EmrStepSensor( task_id='test_task', poke_interval=1, job_flow_id='j-8989898989', step_id='s-VK57YR1Z9Z5N', aws_conn_id='aws_default', ) mock_emr_session = MagicMock() mock_emr_session.client.return_value = self.emr_client_mock # Mock out the emr_client creator self.boto3_session_mock = MagicMock(return_value=mock_emr_session) def test_step_completed(self): self.emr_client_mock.describe_step.side_effect = [ DESCRIBE_JOB_STEP_RUNNING_RETURN, DESCRIBE_JOB_STEP_COMPLETED_RETURN ] with patch('boto3.session.Session', self.boto3_session_mock): self.sensor.execute(None) self.assertEqual(self.emr_client_mock.describe_step.call_count, 2) self.emr_client_mock.describe_step.assert_called_with( ClusterId='j-8989898989', StepId='s-VK57YR1Z9Z5N' ) def test_step_cancelled(self): self.emr_client_mock.describe_step.side_effect = [ DESCRIBE_JOB_STEP_RUNNING_RETURN, DESCRIBE_JOB_STEP_CANCELLED_RETURN ] self.boto3_client_mock = MagicMock(return_value=self.emr_client_mock) with patch('boto3.session.Session', self.boto3_session_mock): self.assertRaises(AirflowException, self.sensor.execute, None) def test_step_interrupted(self): self.emr_client_mock.describe_step.side_effect = [ DESCRIBE_JOB_STEP_RUNNING_RETURN, DESCRIBE_JOB_STEP_INTERRUPTED_RETURN ] self.boto3_client_mock = MagicMock(return_value=self.emr_client_mock) with patch('boto3.session.Session', self.boto3_session_mock): self.assertRaises(AirflowException, self.sensor.execute, None)
def test_execute_calls_with_the_job_flow_id_and_step_id_until_it_reaches_a_terminal_state(self): with patch('boto3.client', self.boto3_client_mock): operator = EmrStepSensor( task_id='test_task', poke_interval=1, job_flow_id='j-8989898989', step_id='s-VK57YR1Z9Z5N', aws_conn_id='aws_default', ) operator.execute(None) # make sure we called twice self.assertEqual(self.mock_emr_client.describe_step.call_count, 2) # make sure it was called with the job_flow_id and step_id self.mock_emr_client.describe_step.assert_called_with(ClusterId='j-8989898989', StepId='s-VK57YR1Z9Z5N')
def test_execute_calls_with_the_job_flow_id_and_step_id_until_it_reaches_a_terminal_state( self): with patch('boto3.client', self.boto3_client_mock): operator = EmrStepSensor( task_id='test_task', poke_interval=1, job_flow_id='j-8989898989', step_id='s-VK57YR1Z9Z5N', aws_conn_id='aws_default', ) operator.execute(None) # make sure we called twice self.assertEqual(self.mock_emr_client.describe_step.call_count, 2) # make sure it was called with the job_flow_id and step_id self.mock_emr_client.describe_step.assert_called_with( ClusterId='j-8989898989', StepId='s-VK57YR1Z9Z5N')
class TestEmrStepSensor(unittest.TestCase): def setUp(self): configuration.load_test_config() self.emr_client_mock = MagicMock() self.sensor = EmrStepSensor( task_id='test_task', poke_interval=1, job_flow_id='j-8989898989', step_id='s-VK57YR1Z9Z5N', aws_conn_id='aws_default', ) def test_step_completed(self): self.emr_client_mock.describe_step.side_effect = [ DESCRIBE_JOB_STEP_RUNNING_RETURN, DESCRIBE_JOB_STEP_COMPLETED_RETURN ] self.boto3_client_mock = MagicMock(return_value=self.emr_client_mock) with patch('boto3.client', self.boto3_client_mock): self.sensor.execute(None) self.assertEqual(self.emr_client_mock.describe_step.call_count, 2) self.emr_client_mock.describe_step.assert_called_with( ClusterId='j-8989898989', StepId='s-VK57YR1Z9Z5N') def test_step_cancelled(self): self.emr_client_mock.describe_step.side_effect = [ DESCRIBE_JOB_STEP_RUNNING_RETURN, DESCRIBE_JOB_STEP_CANCELLED_RETURN ] self.boto3_client_mock = MagicMock(return_value=self.emr_client_mock) with patch('boto3.client', self.boto3_client_mock): self.assertRaises(AirflowException, self.sensor.execute, None)