diff --git a/providers/tests/microsoft/azure/base.py b/providers/tests/microsoft/azure/base.py index 600e4ce488e08..f22212cc3ea12 100644 --- a/providers/tests/microsoft/azure/base.py +++ b/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 diff --git a/providers/tests/microsoft/azure/operators/test_msgraph.py b/providers/tests/microsoft/azure/operators/test_msgraph.py index cd36ae5b131e5..8b0aaf6b82c3e 100644 --- a/providers/tests/microsoft/azure/operators/test_msgraph.py +++ b/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: @@ -60,7 +61,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") @@ -88,7 +89,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 @@ -108,7 +109,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): @@ -128,7 +129,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 @@ -152,7 +153,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 @@ -179,7 +180,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 diff --git a/providers/tests/microsoft/azure/operators/test_powerbi.py b/providers/tests/microsoft/azure/operators/test_powerbi.py index a115b4c52dc50..bf7c03b8c9ee2 100644 --- a/providers/tests/microsoft/azure/operators/test_powerbi.py +++ b/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" diff --git a/providers/tests/microsoft/azure/sensors/test_msgraph.py b/providers/tests/microsoft/azure/sensors/test_msgraph.py index 4240a88e4039b..e3314671ee368 100644 --- a/providers/tests/microsoft/azure/sensors/test_msgraph.py +++ b/providers/tests/microsoft/azure/sensors/test_msgraph.py @@ -26,6 +26,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 @@ -45,7 +46,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) @@ -78,7 +79,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) diff --git a/providers/tests/microsoft/azure/triggers/test_msgraph.py b/providers/tests/microsoft/azure/triggers/test_msgraph.py index ce5a554fe1d52..c80dbe305bf6f 100644 --- a/providers/tests/microsoft/azure/triggers/test_msgraph.py +++ b/providers/tests/microsoft/azure/triggers/test_msgraph.py @@ -42,6 +42,7 @@ mock_json_response, mock_response, ) +from tests_common.test_utils.operators.run_deferrable import run_trigger class TestMSGraphTrigger(Base): @@ -51,7 +52,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) @@ -64,7 +65,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) @@ -75,7 +76,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" @@ -91,7 +92,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" diff --git a/providers/tests/microsoft/conftest.py b/providers/tests/microsoft/conftest.py index d4a65075f5061..e0ec3172cfa2a 100644 --- a/providers/tests/microsoft/conftest.py +++ b/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 @@ -35,15 +34,6 @@ from airflow.models import Connection from airflow.providers.microsoft.azure.hooks.powerbi import PowerBIHook -try: - from airflow.sdk.definitions.context import Context -except ImportError: - # TODO: Remove once provider drops support for Airflow 2 - from airflow.utils.context import Context - -if TYPE_CHECKING: - from sqlalchemy.orm import Session - T = TypeVar("T", dict, str, Connection) @@ -117,52 +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}", 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] - - def remove_license_header(content: str) -> str: """ Removes license header from the given content. diff --git a/tests_common/test_utils/mock_context.py b/tests_common/test_utils/mock_context.py new file mode 100644 index 0000000000000..3c8a101ce2058 --- /dev/null +++ b/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] diff --git a/tests_common/test_utils/operators/run_deferrable.py b/tests_common/test_utils/operators/run_deferrable.py new file mode 100644 index 0000000000000..cdcc781767648 --- /dev/null +++ b/tests_common/test_utils/operators/run_deferrable.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 + +import asyncio +from copy import deepcopy +from typing import TYPE_CHECKING, Any + +from airflow.exceptions import TaskDeferred + +from tests_common.test_utils.mock_context import mock_context + +if TYPE_CHECKING: + from airflow.models import Operator + from airflow.triggers.base import BaseTrigger, TriggerEvent + + +async def run_tigger(trigger: BaseTrigger) -> list[TriggerEvent]: + events = [] + async for event in trigger.run(): + events.append(event) + return events + + +def run_trigger(trigger: BaseTrigger) -> list[TriggerEvent]: + return asyncio.run(run_tigger(trigger)) + + +def execute_operator(operator: Operator) -> tuple[Any, Any]: + context = mock_context(task=operator) + return asyncio.run(deferrable_operator(context, operator)) + + +async def deferrable_operator(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 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