Moved test functions from azure provider to tests_common package for reusability in other providers

providers/tests/microsoft/azure/base.py

@@ -16,22 +16,14 @@
 # under the License.
 from __future__ import annotations
 
-import asyncio
 from contextlib import contextmanager
-from copy import deepcopy
-from typing import TYPE_CHECKING, Any
 from unittest.mock import patch
 
 from kiota_http.httpx_request_adapter import HttpxRequestAdapter
 
-from airflow.exceptions import TaskDeferred
 from airflow.providers.microsoft.azure.hooks.msgraph import KiotaRequestAdapterHook
 
-from providers.tests.microsoft.conftest import get_airflow_connection, mock_context
-
-if TYPE_CHECKING:
-    from airflow.models import Operator
-    from airflow.triggers.base import BaseTrigger, TriggerEvent
+from providers.tests.microsoft.conftest import get_airflow_connection
 
 
 class Base:
@@ -49,42 +41,3 @@ def patch_hook_and_request_adapter(self, response):
             else:
                 mock_get_http_response.return_value = response
             yield
-
-    @staticmethod
-    async def _run_tigger(trigger: BaseTrigger) -> list[TriggerEvent]:
-        events = []
-        async for event in trigger.run():
-            events.append(event)
-        return events
-
-    def run_trigger(self, trigger: BaseTrigger) -> list[TriggerEvent]:
-        return asyncio.run(self._run_tigger(trigger))
-
-    def execute_operator(self, operator: Operator) -> tuple[Any, Any]:
-        context = mock_context(task=operator)
-        return asyncio.run(self.deferrable_operator(context, operator))
-
-    async def deferrable_operator(self, context, operator):
-        result = None
-        triggered_events = []
-        try:
-            operator.render_template_fields(context=context)
-            result = operator.execute(context=context)
-        except TaskDeferred as deferred:
-            task = deferred
-
-            while task:
-                events = await self._run_tigger(task.trigger)
-
-                if not events:
-                    break
-
-                triggered_events.extend(deepcopy(events))
-
-                try:
-                    method = getattr(operator, task.method_name)
-                    result = method(context=context, event=next(iter(events)).payload)
-                    task = None
-                except TaskDeferred as exception:
-                    task = exception
-        return result, triggered_events

providers/tests/microsoft/azure/operators/test_msgraph.py

@@ -31,10 +31,11 @@
 from providers.tests.microsoft.conftest import (
     load_file,
     load_json,
-    mock_context,
     mock_json_response,
     mock_response,
 )
+from tests_common.test_utils.mock_context import mock_context
+from tests_common.test_utils.operators.run_deferrable import execute_operator
 from tests_common.test_utils.version_compat import AIRFLOW_V_2_10_PLUS
 
 if TYPE_CHECKING:
@@ -56,7 +57,7 @@ def test_execute(self):
                 result_processor=lambda context, result: result.get("value"),
             )
 
-            results, events = self.execute_operator(operator)
+            results, events = execute_operator(operator)
 
             assert len(results) == 30
             assert results == users.get("value") + next_users.get("value")
@@ -84,7 +85,7 @@ def test_execute_when_do_xcom_push_is_false(self):
                 do_xcom_push=False,
             )
 
-            results, events = self.execute_operator(operator)
+            results, events = execute_operator(operator)
 
             assert isinstance(results, dict)
             assert len(events) == 1
@@ -104,7 +105,7 @@ def test_execute_when_an_exception_occurs(self):
             )
 
             with pytest.raises(AirflowException):
-                self.execute_operator(operator)
+                execute_operator(operator)
 
     @pytest.mark.db_test
     def test_execute_when_an_exception_occurs_on_custom_event_handler(self):
@@ -124,7 +125,7 @@ def custom_event_handler(context: Context, event: dict[Any, Any] | None = None):
                 event_handler=custom_event_handler,
             )
 
-            results, events = self.execute_operator(operator)
+            results, events = execute_operator(operator)
 
             assert not results
             assert len(events) == 1
@@ -148,7 +149,7 @@ def test_execute_when_response_is_bytes(self):
                 path_parameters={"drive_id": drive_id},
             )
 
-            results, events = self.execute_operator(operator)
+            results, events = execute_operator(operator)
 
             assert operator.path_parameters == {"drive_id": drive_id}
             assert results == base64_encoded_content
@@ -175,7 +176,7 @@ def test_execute_with_lambda_parameter_when_response_is_bytes(self):
                 path_parameters=lambda context, jinja_env: {"drive_id": drive_id},
             )
 
-            results, events = self.execute_operator(operator)
+            results, events = execute_operator(operator)
 
             assert operator.path_parameters == {"drive_id": drive_id}
             assert results == base64_encoded_content

providers/tests/microsoft/azure/operators/test_powerbi.py

@@ -32,7 +32,8 @@
 from airflow.utils import timezone
 
 from providers.tests.microsoft.azure.base import Base
-from providers.tests.microsoft.conftest import get_airflow_connection, mock_context
+from providers.tests.microsoft.conftest import get_airflow_connection
+from tests_common.test_utils.mock_context import mock_context
 
 DEFAULT_CONNECTION_CLIENT_SECRET = "powerbi_conn_id"
 TASK_ID = "run_powerbi_operator"

providers/tests/microsoft/azure/sensors/test_msgraph.py

@@ -25,6 +25,7 @@
 
 from providers.tests.microsoft.azure.base import Base
 from providers.tests.microsoft.conftest import load_json, mock_json_response
+from tests_common.test_utils.operators.run_deferrable import execute_operator
 from tests_common.test_utils.version_compat import AIRFLOW_V_2_10_PLUS
 
 
@@ -43,7 +44,7 @@ def test_execute(self):
                 timeout=350.0,
             )
 
-            results, events = self.execute_operator(sensor)
+            results, events = execute_operator(sensor)
 
             assert sensor.path_parameters == {"scanId": "0a1b1bf3-37de-48f7-9863-ed4cda97a9ef"}
             assert isinstance(results, str)
@@ -69,7 +70,7 @@ def test_execute_with_lambda_parameter(self):
                 timeout=350.0,
             )
 
-            results, events = self.execute_operator(sensor)
+            results, events = execute_operator(sensor)
 
             assert sensor.path_parameters == {"scanId": "0a1b1bf3-37de-48f7-9863-ed4cda97a9ef"}
             assert isinstance(results, str)

providers/tests/microsoft/azure/triggers/test_msgraph.py

@@ -40,6 +40,7 @@
     mock_json_response,
     mock_response,
 )
+from tests_common.test_utils.operators.run_deferrable import run_trigger
 
 
 class TestMSGraphTrigger(Base):
@@ -49,7 +50,7 @@ def test_run_when_valid_response(self):
 
         with self.patch_hook_and_request_adapter(response):
             trigger = MSGraphTrigger("users/delta", conn_id="msgraph_api")
-            actual = self.run_trigger(trigger)
+            actual = run_trigger(trigger)
 
             assert len(actual) == 1
             assert isinstance(actual[0], TriggerEvent)
@@ -62,7 +63,7 @@ def test_run_when_response_is_none(self):
 
         with self.patch_hook_and_request_adapter(response):
             trigger = MSGraphTrigger("users/delta", conn_id="msgraph_api")
-            actual = self.run_trigger(trigger)
+            actual = run_trigger(trigger)
 
             assert len(actual) == 1
             assert isinstance(actual[0], TriggerEvent)
@@ -73,7 +74,7 @@ def test_run_when_response_is_none(self):
     def test_run_when_response_cannot_be_converted_to_json(self):
         with self.patch_hook_and_request_adapter(AirflowException()):
             trigger = MSGraphTrigger("users/delta", conn_id="msgraph_api")
-            actual = next(iter(self.run_trigger(trigger)))
+            actual = next(iter(run_trigger(trigger)))
 
             assert isinstance(actual, TriggerEvent)
             assert actual.payload["status"] == "failure"
@@ -89,7 +90,7 @@ def test_run_when_response_is_bytes(self):
                 "https://graph.microsoft.com/v1.0/me/drive/items/1b30fecf-4330-4899-b249-104c2afaf9ed/content"
             )
             trigger = MSGraphTrigger(url, response_type="bytes", conn_id="msgraph_api")
-            actual = next(iter(self.run_trigger(trigger)))
+            actual = next(iter(run_trigger(trigger)))
 
             assert isinstance(actual, TriggerEvent)
             assert actual.payload["status"] == "success"

providers/tests/microsoft/conftest.py

@@ -21,11 +21,10 @@
 import random
 import re
 import string
-from collections.abc import Iterable
 from inspect import currentframe
 from json import JSONDecodeError
 from os.path import dirname, join
-from typing import TYPE_CHECKING, Any, TypeVar
+from typing import Any, TypeVar
 from unittest.mock import MagicMock
 
 import pytest
@@ -34,10 +33,6 @@
 
 from airflow.models import Connection
 from airflow.providers.microsoft.azure.hooks.powerbi import PowerBIHook
-from airflow.utils.context import Context
-
-if TYPE_CHECKING:
-    from sqlalchemy.orm import Session
 
 T = TypeVar("T", dict, str, Connection)
 
@@ -112,50 +107,6 @@ def mock_response(status_code, content: Any = None, headers: dict | None = None)
     return response
 
 
-def mock_context(task) -> Context:
-    from airflow.models import TaskInstance
-    from airflow.utils.session import NEW_SESSION
-    from airflow.utils.state import TaskInstanceState
-    from airflow.utils.xcom import XCOM_RETURN_KEY
-
-    values: dict[str, Any] = {}
-
-    class MockedTaskInstance(TaskInstance):
-        def __init__(
-            self,
-            task,
-            run_id: str | None = "run_id",
-            state: str | None = TaskInstanceState.RUNNING,
-            map_index: int = -1,
-        ):
-            super().__init__(task=task, run_id=run_id, state=state, map_index=map_index)
-            self.values: dict[str, Any] = {}
-
-        def xcom_pull(
-            self,
-            task_ids: str | Iterable[str] | None = None,
-            dag_id: str | None = None,
-            key: str = XCOM_RETURN_KEY,
-            include_prior_dates: bool = False,
-            session: Session = NEW_SESSION,
-            *,
-            map_indexes: int | Iterable[int] | None = None,
-            default: Any = None,
-            run_id: str | None = None,
-        ) -> Any:
-            if map_indexes:
-                return values.get(f"{task_ids or self.task_id}_{dag_id or self.dag_id}_{key}_{map_indexes}")
-            return values.get(f"{task_ids or self.task_id}_{dag_id or self.dag_id}_{key}")
-
-        def xcom_push(self, key: str, value: Any, session: Session = NEW_SESSION, **kwargs) -> None:
-            values[f"{self.task_id}_{self.dag_id}_{key}_{self.map_index}"] = value
-
-    values["ti"] = MockedTaskInstance(task=task)
-
-    # See https://github.com/python/mypy/issues/8890 - mypy does not support passing typed dict to TypedDict
-    return Context(values)  # type: ignore[misc]
-
-
 def remove_license_header(content: str) -> str:
     """
     Removes license header from the given content.

tests_common/test_utils/mock_context.py

@@ -0,0 +1,71 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from collections.abc import Iterable
+from typing import TYPE_CHECKING, Any
+
+from airflow.utils.context import Context
+
+if TYPE_CHECKING:
+    from sqlalchemy.orm import Session
+
+
+def mock_context(task) -> Context:
+    from airflow.models import TaskInstance
+    from airflow.utils.session import NEW_SESSION
+    from airflow.utils.state import TaskInstanceState
+    from airflow.utils.xcom import XCOM_RETURN_KEY
+
+    values: dict[str, Any] = {}
+
+    class MockedTaskInstance(TaskInstance):
+        def __init__(
+            self,
+            task,
+            run_id: str | None = "run_id",
+            state: str | None = TaskInstanceState.RUNNING,
+            map_index: int = -1,
+        ):
+            super().__init__(task=task, run_id=run_id, state=state, map_index=map_index)
+            self.values: dict[str, Any] = {}
+
+        def xcom_pull(
+            self,
+            task_ids: str | Iterable[str] | None = None,
+            dag_id: str | None = None,
+            key: str = XCOM_RETURN_KEY,
+            include_prior_dates: bool = False,
+            session: Session = NEW_SESSION,
+            *,
+            map_indexes: int | Iterable[int] | None = None,
+            default: Any = None,
+            run_id: str | None = None,
+        ) -> Any:
+            if map_indexes:
+                return values.get(
+                    f"{task_ids or self.task_id}_{dag_id or self.dag_id}_{key}_{map_indexes}", default
+                )
+            return values.get(f"{task_ids or self.task_id}_{dag_id or self.dag_id}_{key}", default)
+
+        def xcom_push(self, key: str, value: Any, session: Session = NEW_SESSION, **kwargs) -> None:
+            values[f"{self.task_id}_{self.dag_id}_{key}_{self.map_index}"] = value
+
+    values["ti"] = MockedTaskInstance(task=task)
+
+    # See https://github.com/python/mypy/issues/8890 - mypy does not support passing typed dict to TypedDict
+    return Context(values)  # type: ignore[misc]