From 36eba1885a7e0e720dadef570b41a741a717d275 Mon Sep 17 00:00:00 2001
From: Pankaj Singh <98807258+pankajastro@users.noreply.github.com>
Date: Mon, 11 Sep 2023 23:06:32 +0530
Subject: [PATCH] Fix Databricks trigger test (#34272)

---
 .../databricks/triggers/test_databricks.py    | 45 +++++++++----------
 1 file changed, 22 insertions(+), 23 deletions(-)

diff --git a/tests/providers/databricks/triggers/test_databricks.py b/tests/providers/databricks/triggers/test_databricks.py
index 53f67eb251833..91dcbc418b58e 100644
--- 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)