changes after CR apache#2

aviemzur · Dec 31, 2019 · d1d638f · d1d638f
1 parent 8bf40f9
commit d1d638f
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 6 deletions.
diff --git a/airflow/contrib/operators/emr_add_steps_operator.py b/airflow/contrib/operators/emr_add_steps_operator.py
@@ -72,12 +72,10 @@ def execute(self, context):
 
         emr = emr_hook.get_conn()
 
-        job_flow_id = self.job_flow_id
-
+        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_hook.get_cluster_id_by_name(self.job_flow_name, self.cluster_states)
-            if not job_flow_id:
-                raise AirflowException(f'No cluster found for name: {self.job_flow_name}')
+            raise AirflowException(f'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
@@ -146,7 +146,7 @@ def test_init_with_nonexistent_cluster_name(self):
 
             with self.assertRaises(AirflowException) as error:
                 operator.execute(self.mock_context)
-            self.assertEqual(str(error.exception), f'No cluster found for name:  {cluster_name}')
+            self.assertEqual(str(error.exception), f'No cluster found for name: {cluster_name}')
 
 
 if __name__ == '__main__':