Bug fix for EmrAddStepOperator init with cluster_name error (apache#9235

) Closes apache#9127
cfei18 · Mar 5, 2021 · 3f4db1c · 3f4db1c
1 parent 7cbb758
commit 3f4db1c
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 4 deletions.
diff --git a/airflow/contrib/operators/emr_add_steps_operator.py b/airflow/contrib/operators/emr_add_steps_operator.py
@@ -66,12 +66,14 @@ def __init__(
         self.steps = steps
 
     def execute(self, context):
-        emr = EmrHook(aws_conn_id=self.aws_conn_id).get_conn()
+        emr_hook = EmrHook(aws_conn_id=self.aws_conn_id)
 
-        job_flow_id = self.job_flow_id
+        emr = emr_hook.get_conn()
 
+        job_flow_id = self.job_flow_id or emr_hook.get_cluster_id_by_name(self.job_flow_name,
+                                                                          self.cluster_states)
         if not job_flow_id:
-            job_flow_id = emr.get_cluster_id_by_name(self.job_flow_name, self.cluster_states)
+            raise AirflowException('No cluster found for name: ' + self.job_flow_name)
 
         if self.do_xcom_push:
             context['ti'].xcom_push(key='job_flow_id', value=job_flow_id)

diff --git a/tests/contrib/operators/test_emr_add_steps_operator.py b/tests/contrib/operators/test_emr_add_steps_operator.py
@@ -105,10 +105,11 @@ def test_execute_returns_step_id(self):
         with patch('boto3.session.Session', self.boto3_session_mock):
             self.assertEqual(self.operator.execute(self.mock_context), ['s-2LH3R5GW3A53T'])
 
+    @patch.multiple('airflow.contrib.hooks.emr_hook.EmrHook',
+                    get_cluster_id_by_name=MagicMock(return_value='j-1231231234'))
     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):