Fix EmrNotebookExecutionSensor broken test (#38354)

apache · Mar 21, 2024 · 952d04a · 952d04a
1 parent 3400dfe
commit 952d04a
Showing 1 changed file with 5 additions and 4 deletions.
diff --git a/tests/providers/amazon/aws/sensors/test_emr_notebook_execution.py b/tests/providers/amazon/aws/sensors/test_emr_notebook_execution.py
@@ -65,12 +65,13 @@ def test_emr_notebook_execution_sensor_failed_state(self, mock_conn):
     @mock.patch("airflow.providers.amazon.aws.hooks.emr.EmrHook.conn")
     def test_emr_notebook_execution_sensor_success_state_multiple(self, mock_conn):
         return_values = [self._generate_response("PENDING") for i in range(2)]
-        return_values = return_values.append(self._generate_response("FINISHED"))
-        mock_conn.describe_notebook_execution.side_effects = return_values
+        return_values.append(self._generate_response("FINISHED"))
+        mock_conn.describe_notebook_execution.side_effect = return_values
         sensor = EmrNotebookExecutionSensor(
             task_id="test_task",
             poke_interval=0,
             notebook_execution_id="test-execution-id",
         )
-        sensor.poke(None)
-        mock_conn.describe_notebook_execution.call_count == 3
+        while not sensor.poke(None):
+            pass
+        assert mock_conn.describe_notebook_execution.call_count == 3