Fix Databricks trigger test (#34272)

apache · Sep 11, 2023 · 36eba18 · 36eba18
1 parent d3ce442
commit 36eba18
Showing 1 changed file with 22 additions and 23 deletions.
diff --git a/tests/providers/databricks/triggers/test_databricks.py b/tests/providers/databricks/triggers/test_databricks.py
@@ -107,23 +107,22 @@ async def test_run_return_success(self, mock_get_run_state, mock_get_run_page_ur
             result_state="SUCCESS",
         )
 
-        generator = self.trigger.run()
-        actual = await generator.asend(None)
-        assert actual == TriggerEvent(
-            {
-                "run_id": RUN_ID,
-                "run_state": RunState(
-                    life_cycle_state=LIFE_CYCLE_STATE_TERMINATED, state_message="", result_state="SUCCESS"
-                ).to_json(),
-                "run_page_url": RUN_PAGE_URL,
-            }
-        )
+        trigger_event = self.trigger.run()
+        async for event in trigger_event:
+            assert event == TriggerEvent(
+                {
+                    "run_id": RUN_ID,
+                    "run_state": RunState(
+                        life_cycle_state=LIFE_CYCLE_STATE_TERMINATED, state_message="", result_state="SUCCESS"
+                    ).to_json(),
+                    "run_page_url": RUN_PAGE_URL,
+                }
+            )
 
     @pytest.mark.asyncio
     @mock.patch("airflow.providers.databricks.triggers.databricks.asyncio.sleep")
     @mock.patch("airflow.providers.databricks.hooks.databricks.DatabricksHook.a_get_run_state")
     async def test_sleep_between_retries(self, mock_get_run_state, mock_sleep):
-
         mock_get_run_state.side_effect = [
             RunState(
                 life_cycle_state=LIFE_CYCLE_STATE_PENDING,
@@ -137,16 +136,16 @@ async def test_sleep_between_retries(self, mock_get_run_state, mock_sleep):
             ),
         ]
 
-        generator = self.trigger.run()
-        actual = await generator.asend(None)
-        assert actual == TriggerEvent(
-            {
-                "run_id": RUN_ID,
-                "run_state": RunState(
-                    life_cycle_state=LIFE_CYCLE_STATE_TERMINATED, state_message="", result_state="SUCCESS"
-                ).to_json(),
-                "run_page_url": RUN_PAGE_URL,
-            }
-        )
+        trigger_event = self.trigger.run()
+        async for event in trigger_event:
+            assert event == TriggerEvent(
+                {
+                    "run_id": RUN_ID,
+                    "run_state": RunState(
+                        life_cycle_state=LIFE_CYCLE_STATE_TERMINATED, state_message="", result_state="SUCCESS"
+                    ).to_json(),
+                    "run_page_url": RUN_PAGE_URL,
+                }
+            )
         mock_sleep.assert_called_once()
         mock_sleep.assert_called_with(POLLING_INTERVAL_SECONDS)