def test_execute_calls_with_the_job_flow_id_until_it_reaches_a_terminal_state(self):
with patch('boto3.session.Session', self.boto3_session_mock):
operator = EmrJobFlowSensor(
task_id='test_task',
poke_interval=2,
job_flow_id='j-8989898989',
aws_conn_id='aws_default'
)
operator.execute(None)
# make sure we called twice
self.assertEqual(self.mock_emr_client.describe_cluster.call_count, 2)
# make sure it was called with the job_flow_id
self.mock_emr_client.describe_cluster.assert_called_with(ClusterId='j-8989898989')
def test_execute_calls_with_the_job_flow_id_until_it_reaches_a_terminal_state(self):
with patch('boto3.session.Session', self.boto3_session_mock):
operator = EmrJobFlowSensor(
task_id='test_task',
poke_interval=2,
job_flow_id='j-8989898989',
aws_conn_id='aws_default'
)
operator.execute(None)
# make sure we called twice
self.assertEqual(self.mock_emr_client.describe_cluster.call_count, 2)
# make sure it was called with the job_flow_id
self.mock_emr_client.describe_cluster.assert_called_with(ClusterId='j-8989898989')
def test_execute_calls_with_the_job_flow_id_until_it_reaches_failed_state_with_exception(self):
self.mock_emr_client.describe_cluster.side_effect = [
DESCRIBE_CLUSTER_RUNNING_RETURN,
DESCRIBE_CLUSTER_TERMINATED_WITH_ERRORS_RETURN
]
with patch('boto3.session.Session', self.boto3_session_mock):
operator = EmrJobFlowSensor(
task_id='test_task',
poke_interval=2,
job_flow_id='j-8989898989',
aws_conn_id='aws_default'
)
with self.assertRaises(AirflowException):
operator.execute(None)
# make sure we called twice
self.assertEqual(self.mock_emr_client.describe_cluster.call_count, 2)
# make sure it was called with the job_flow_id
self.mock_emr_client.describe_cluster.assert_called_once_with(ClusterId='j-8989898989')
def test_execute_calls_with_the_job_flow_id_until_it_reaches_a_terminal_state(
self):
self.mock_emr_client.describe_cluster.side_effect = [
DESCRIBE_CLUSTER_RUNNING_RETURN, DESCRIBE_CLUSTER_TERMINATED_RETURN
]
with patch('boto3.session.Session', self.boto3_session_mock):
operator = EmrJobFlowSensor(task_id='test_task',
poke_interval=2,
job_flow_id='j-8989898989',
aws_conn_id='aws_default')
operator.execute(None)
# make sure we called twice
self.assertEqual(self.mock_emr_client.describe_cluster.call_count,
2)
# make sure it was called with the job_flow_id
calls = [
unittest.mock.call(ClusterId='j-8989898989'),
unittest.mock.call(ClusterId='j-8989898989')
]
self.mock_emr_client.describe_cluster.assert_has_calls(calls)
