def test_different_target_states(self):
self.mock_emr_client.describe_cluster.side_effect = [
DESCRIBE_CLUSTER_STARTING_RETURN, # return False
DESCRIBE_CLUSTER_BOOTSTRAPPING_RETURN, # return False
DESCRIBE_CLUSTER_RUNNING_RETURN, # return True
DESCRIBE_CLUSTER_WAITING_RETURN, # will not be used
DESCRIBE_CLUSTER_TERMINATED_RETURN, # will not be used
DESCRIBE_CLUSTER_TERMINATED_WITH_ERRORS_RETURN, # will not be used
]
with patch('boto3.session.Session', self.boto3_session_mock):
operator = EmrJobFlowSensor(
task_id='test_task',
poke_interval=0,
job_flow_id='j-8989898989',
aws_conn_id='aws_default',
target_states=['RUNNING', 'WAITING'],
)
operator.execute(None)
# make sure we called twice
assert self.mock_emr_client.describe_cluster.call_count == 3
# make sure it was called with the job_flow_id
calls = [unittest.mock.call(ClusterId='j-8989898989')]
self.mock_emr_client.describe_cluster.assert_has_calls(calls)
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=0,
job_flow_id='j-8989898989',
aws_conn_id='aws_default')
with pytest.raises(AirflowException):
operator.execute(None)
# make sure we called twice
assert 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_target_state(
self):
self.mock_emr_client.describe_cluster.side_effect = [
DESCRIBE_CLUSTER_STARTING_RETURN,
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=0,
job_flow_id='j-8989898989',
aws_conn_id='aws_default')
operator.execute(None)
# make sure we called twice
assert self.mock_emr_client.describe_cluster.call_count == 3
# make sure it was called with the job_flow_id
calls = [unittest.mock.call(ClusterId='j-8989898989')]
self.mock_emr_client.describe_cluster.assert_has_calls(calls)