From b9513cc9b94a70e721a439088ddd9e93f8c6c429 Mon Sep 17 00:00:00 2001
From: ambikagarg <70703123+ambika-garg@users.noreply.github.com>
Date: Sat, 9 Nov 2024 13:09:50 -0500
Subject: [PATCH] Fix Power BI trigger testcase (#43494)

* Optimize test that check trigger after being timeout

* Fix typing error

* Update test_powerbi.py

* refactor: Refactored powerbi_trigger fixture and used pytest.mark.parametrize to parametrize timeout

* Refactor: Remove parameterize marker

---------

Co-authored-by: David Blain <info@dabla.be>
---
 .../microsoft/azure/triggers/powerbi.py       | 19 ++++++++++++-------
 .../microsoft/azure/triggers/test_powerbi.py  | 19 +++++++++----------
 2 files changed, 21 insertions(+), 17 deletions(-)

diff --git a/providers/src/airflow/providers/microsoft/azure/triggers/powerbi.py b/providers/src/airflow/providers/microsoft/azure/triggers/powerbi.py
index a74898f55f28..3e23f3fa068d 100644
--- a/providers/src/airflow/providers/microsoft/azure/triggers/powerbi.py
+++ b/providers/src/airflow/providers/microsoft/azure/triggers/powerbi.py
@@ -104,16 +104,21 @@ async def run(self) -> AsyncIterator[TriggerEvent]:
             dataset_id=self.dataset_id,
             group_id=self.group_id,
         )
+
+        async def fetch_refresh_status() -> str:
+            """Fetch the current status of the dataset refresh."""
+            refresh_details = await self.hook.get_refresh_details_by_refresh_id(
+                dataset_id=self.dataset_id,
+                group_id=self.group_id,
+                refresh_id=self.dataset_refresh_id,
+            )
+            return refresh_details["status"]
+
         try:
-            dataset_refresh_status = None
+            dataset_refresh_status = await fetch_refresh_status()
             start_time = time.monotonic()
             while start_time + self.timeout > time.monotonic():
-                refresh_details = await self.hook.get_refresh_details_by_refresh_id(
-                    dataset_id=self.dataset_id,
-                    group_id=self.group_id,
-                    refresh_id=self.dataset_refresh_id,
-                )
-                dataset_refresh_status = refresh_details.get("status")
+                dataset_refresh_status = await fetch_refresh_status()
 
                 if dataset_refresh_status == PowerBIDatasetRefreshStatus.COMPLETED:
                     yield TriggerEvent(
diff --git a/providers/tests/microsoft/azure/triggers/test_powerbi.py b/providers/tests/microsoft/azure/triggers/test_powerbi.py
index 8f5a3e945ffc..edd1a60e180e 100644
--- a/providers/tests/microsoft/azure/triggers/test_powerbi.py
+++ b/providers/tests/microsoft/azure/triggers/test_powerbi.py
@@ -32,27 +32,26 @@
 DATASET_ID = "dataset_id"
 GROUP_ID = "group_id"
 DATASET_REFRESH_ID = "dataset_refresh_id"
-TIMEOUT = 30
+TIMEOUT = 5
 MODULE = "airflow.providers.microsoft.azure"
-CHECK_INTERVAL = 10
+CHECK_INTERVAL = 1
 API_VERSION = "v1.0"
 
 
 @pytest.fixture
-def powerbi_trigger():
-    trigger = PowerBITrigger(
+def powerbi_trigger(timeout=TIMEOUT, check_interval=CHECK_INTERVAL) -> PowerBITrigger:
+    """Fixture for creating a PowerBITrigger with customizable timeout and check interval."""
+    return PowerBITrigger(
         conn_id=POWERBI_CONN_ID,
         proxies=None,
         api_version=API_VERSION,
         dataset_id=DATASET_ID,
         group_id=GROUP_ID,
-        check_interval=CHECK_INTERVAL,
+        check_interval=check_interval,
         wait_for_termination=True,
-        timeout=TIMEOUT,
+        timeout=timeout,
     )
 
-    return trigger
-
 
 class TestPowerBITrigger:
     @mock.patch("airflow.hooks.base.BaseHook.get_connection", side_effect=get_airflow_connection)
@@ -232,7 +231,7 @@ async def test_powerbi_trigger_run_exception_without_refresh_id(
     async def test_powerbi_trigger_run_timeout(
         self, mock_trigger_dataset_refresh, mock_get_refresh_details_by_refresh_id, powerbi_trigger
     ):
-        """Assert that powerbi run timesout after end_time elapses"""
+        """Assert that powerbi run times out after end_time elapses"""
         mock_get_refresh_details_by_refresh_id.return_value = {
             "status": PowerBIDatasetRefreshStatus.IN_PROGRESS
         }
@@ -243,7 +242,7 @@ async def test_powerbi_trigger_run_timeout(
         expected = TriggerEvent(
             {
                 "status": "error",
-                "message": f"Timeout occurred while waiting for dataset refresh to complete: The dataset refresh {DATASET_REFRESH_ID} has status In Progress.",
+                "message": f"Timeout occurred while waiting for dataset refresh to complete: The dataset refresh {DATASET_REFRESH_ID} has status {PowerBIDatasetRefreshStatus.IN_PROGRESS}.",
                 "dataset_refresh_id": DATASET_REFRESH_ID,
             }
         )