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)
