ashb commented on a change in pull request #6655: [AIRFLOW-6056] Allow
EmrAddStepsOperator to accept job_flow_name as alternative to job_flow_id
URL: https://github.com/apache/airflow/pull/6655#discussion_r353696071
##########
File path: tests/contrib/operators/test_emr_add_steps_operator.py
##########
@@ -93,13 +101,30 @@ def test_render_template(self):
def test_execute_returns_step_id(self):
self.emr_client_mock.add_job_flow_steps.return_value =
ADD_STEPS_SUCCESS_RETURN
- # Mock out the emr_client creator
- emr_session_mock = MagicMock()
- emr_session_mock.client.return_value = self.emr_client_mock
- self.boto3_session_mock = MagicMock(return_value=emr_session_mock)
+ with patch('boto3.session.Session', self.boto3_session_mock):
+ self.assertEqual(self.operator.execute(self.mock_context),
['s-2LH3R5GW3A53T'])
+
+ def test_init_with_cluster_name(self):
+ expected_job_flow_id = 'j-1231231234'
+
+ self.emr_client_mock.get_cluster_id_by_name.return_value =
expected_job_flow_id
+ self.emr_client_mock.add_job_flow_steps.return_value =
ADD_STEPS_SUCCESS_RETURN
with patch('boto3.session.Session', self.boto3_session_mock):
- self.assertEqual(self.operator.execute(None), ['s-2LH3R5GW3A53T'])
+ operator = EmrAddStepsOperator(
+ task_id='test_task',
+ job_flow_name='test_cluster',
+ cluster_states=['RUNNING', 'WAITING'],
+ aws_conn_id='aws_default',
+ dag=DAG('test_dag_id', default_args=self.args)
+ )
+
+ operator.execute(self.mock_context)
+
+ ti = self.mock_context['ti']
+ job_flow_id = ti.xcom_push.mock_calls[0][2]['value']
+
+ self.assertEqual(job_flow_id, expected_job_flow_id)
Review comment:
This would be more clear as:
```python
ti.xcom_push.assert_any_call(key='job_flow_id',
value=expected_job_flow_id)
```
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
With regards,
Apache Git Services
