This is an automated email from the ASF dual-hosted git repository.
shahar1 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new e74e2584cab Let MSGraphAsyncOperator and MSGraphSensor start from the
triggerer (#53009)
e74e2584cab is described below
commit e74e2584cab6251741d72f962cb014d1ad316e48
Author: David Blain <[email protected]>
AuthorDate: Fri Nov 13 07:03:47 2026 +0100
Let MSGraphAsyncOperator and MSGraphSensor start from the triggerer (#53009)
* Let MSGraphAsyncOperator and MSGraphSensor start from the triggerer
Both do all their requests in the triggerer, so the worker they start on
only hands the first request over and that round trip is wasted. Since
Airflow 3.3 the triggerer renders the templated fields of a task which
starts from a trigger, which was what kept these two from using it.
It is opt-in because the triggerer renders with a plain Jinja environment,
without the user defined macros and filters or the native rendering of the
Dag, and the task no longer waits for a pool slot. A task falls back to a
worker when it cannot start from the trigger: on older Airflow versions,
when an argument is an XComArg, a callable or a file-like object, which
the triggerer cannot resolve or receive, and when it is mapped.
* Keep tasks started from the triggerer consistent with retries, the sensor
timeout and rendered values
- A durable task starts on a worker, as its crash recovery runs in execute.
- Starting from the triggerer needs Airflow 3.4: before, a task the
scheduler
deferred keeps the XComs of an earlier try, and the operator reads them
back.
- The operator and the sensor build their first trigger in one method, and
the
start arguments come from that trigger, so both paths send the same
request
and the sensor no longer passes its timeout as HTTP timeout.
- The sensor counts its timeout from the start of the try when the first
poll
comes from the triggerer.
- MSGraphTrigger takes a rendered flag: triggers deferred by a worker hold
rendered values or values from a response, and the triggerer leaves them
as
they are instead of rendering them again.
---
.../test_utils/operators/run_deferrable.py | 32 ++-
.../microsoft/azure/docs/operators/msgraph.rst | 35 +++
providers/microsoft/azure/docs/sensors/msgraph.rst | 19 ++
.../providers/microsoft/azure/operators/msgraph.py | 102 ++++++--
.../providers/microsoft/azure/sensors/msgraph.py | 67 +++--
.../providers/microsoft/azure/triggers/msgraph.py | 14 +
.../providers/microsoft/azure/version_compat.py | 2 +
.../unit/microsoft/azure/operators/test_msgraph.py | 287 ++++++++++++++++++++-
.../unit/microsoft/azure/sensors/test_msgraph.py | 201 ++++++++++++++-
.../unit/microsoft/azure/triggers/test_msgraph.py | 14 +
10 files changed, 728 insertions(+), 45 deletions(-)
diff --git
a/devel-common/src/tests_common/test_utils/operators/run_deferrable.py
b/devel-common/src/tests_common/test_utils/operators/run_deferrable.py
index cdcc7817676..10919a1f2fe 100644
--- a/devel-common/src/tests_common/test_utils/operators/run_deferrable.py
+++ b/devel-common/src/tests_common/test_utils/operators/run_deferrable.py
@@ -18,6 +18,7 @@ from __future__ import annotations
import asyncio
from copy import deepcopy
+from importlib import import_module
from typing import TYPE_CHECKING, Any
from airflow.exceptions import TaskDeferred
@@ -45,14 +46,43 @@ def execute_operator(operator: Operator) -> tuple[Any, Any]:
return asyncio.run(deferrable_operator(context, operator))
+def start_from_trigger(context, operator) -> TaskDeferred:
+ """
+ Defer an operator which starts from the triggerer, as the scheduler and
the triggerer do.
+
+ The trigger is built from the arguments the operator declared when it was
created, so before
+ any of its templated fields got rendered, and the triggerer then renders
the fields the
+ trigger shares with the operator.
+ """
+ start_trigger_args = operator.start_trigger_args
+ trigger_kwargs = start_trigger_args.trigger_kwargs or {}
+ module_name, _, class_name = start_trigger_args.trigger_cls.rpartition(".")
+ trigger = getattr(import_module(module_name), class_name)(**trigger_kwargs)
+ trigger.template_fields = tuple(
+ field for field in operator.template_fields if field in trigger_kwargs
and hasattr(trigger, field)
+ )
+ trigger.render_template_fields(context=context)
+ return TaskDeferred(
+ trigger=trigger,
+ method_name=start_trigger_args.next_method,
+ kwargs=start_trigger_args.next_kwargs,
+ timeout=start_trigger_args.timeout,
+ )
+
+
async def deferrable_operator(context, operator):
result = None
triggered_events = []
try:
+ if getattr(operator, "start_from_trigger", False):
+ deferred = start_from_trigger(context, operator)
+ # The worker only gets involved once the trigger fired, and
renders the operator then.
+ operator.render_template_fields(context=context)
+ raise deferred
operator.render_template_fields(context=context)
result = operator.execute(context=context)
except TaskDeferred as deferred:
- task = deferred
+ task: TaskDeferred | None = deferred
while task:
events = await run_tigger(task.trigger)
diff --git a/providers/microsoft/azure/docs/operators/msgraph.rst
b/providers/microsoft/azure/docs/operators/msgraph.rst
index 875d3df4772..67fc9ed0c59 100644
--- a/providers/microsoft/azure/docs/operators/msgraph.rst
+++ b/providers/microsoft/azure/docs/operators/msgraph.rst
@@ -95,6 +95,41 @@ from the first page.
Pages are stored as they are, so pages larger than ``[state_store]
max_value_storage_bytes`` need a state
store backend that offloads large values. Durable execution requires Airflow
3.3 or newer.
+.. _howto/operator:MSGraphAsyncOperator:start_from_trigger:
+
+Starting from the triggerer
+~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+The operator does all its requests in the triggerer, so the only thing it does
on a worker at the start
+is to hand the first request over. With ``start_from_trigger=True`` the
scheduler defers the task itself
+and that first worker run is skipped. The worker gets involved once the
response is there, to process it
+and to follow the pagination.
+
+.. code-block:: python
+
+ site_task = MSGraphAsyncOperator(
+ task_id="news_site",
+ conn_id="msgraph_api",
+ url="sites/{{ params.site }}",
+ start_from_trigger=True,
+ )
+
+This needs Airflow 3.4 or later, on older versions the task starts on a worker
as it does without the
+argument. It also starts on a worker when:
+
+* an argument which is passed on to the trigger is an ``XComArg``, a callable
or a file-like object,
+ as the triggerer cannot resolve or receive those;
+* the task is mapped with ``expand``;
+* ``durable=True`` is set, as a retry then resumes from the task state store
on a worker.
+
+Keep the following in mind when you enable it:
+
+* The templated fields of the first request are rendered by the triggerer,
which does not know the
+ ``user_defined_macros`` and ``user_defined_filters`` of the Dag and always
renders strings, also when
+ ``render_template_as_native_obj`` is set.
+* The task goes straight from scheduled to deferred, so the first request does
not wait for a pool
+ slot or for a worker.
+
Cross-References
----------------
diff --git a/providers/microsoft/azure/docs/sensors/msgraph.rst
b/providers/microsoft/azure/docs/sensors/msgraph.rst
index f5477dcd772..64f44a4ef9d 100644
--- a/providers/microsoft/azure/docs/sensors/msgraph.rst
+++ b/providers/microsoft/azure/docs/sensors/msgraph.rst
@@ -33,6 +33,25 @@ Below is an example of using this sensor to poll the status
of a PowerBI workspa
:start-after: [START howto_sensor_powerbi_scan_status]
:end-before: [END howto_sensor_powerbi_scan_status]
+Starting from the triggerer
+^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+With ``start_from_trigger=True`` the scheduler defers the sensor itself, so
the first poll is done without
+a worker being involved. The ``timeout`` of the sensor counts from the moment
the scheduler deferred the
+sensor, so it covers that first poll and the ones after it. The requirements
+and limitations are the same as for the operator, see
+:ref:`howto/operator:MSGraphAsyncOperator:start_from_trigger`.
+
+.. code-block:: python
+
+ check_workspace_status_task = MSGraphSensor(
+ task_id="check_workspaces_status",
+ conn_id="powerbi",
+ url="myorg/admin/workspaces/scanStatus/{scanId}",
+ path_parameters={"scanId": "{{ params.scan_id }}"},
+ start_from_trigger=True,
+ )
+
Reference
---------
diff --git
a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/operators/msgraph.py
b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/operators/msgraph.py
index 12f1f0e1314..4b06ace3209 100644
---
a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/operators/msgraph.py
+++
b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/operators/msgraph.py
@@ -36,9 +36,11 @@ from airflow.providers.microsoft.azure.triggers.msgraph
import (
MSGraphTrigger,
ResponseSerializer,
)
-from airflow.providers.microsoft.azure.version_compat import AIRFLOW_V_3_3_PLUS
+from airflow.providers.microsoft.azure.version_compat import
AIRFLOW_V_3_3_PLUS, AIRFLOW_V_3_4_PLUS
+from airflow.triggers.base import StartTriggerArgs
if TYPE_CHECKING:
+ from datetime import timedelta
from io import BytesIO
from msgraph_core import APIVersion
@@ -64,6 +66,52 @@ class _PaginationCursor(TypedDict):
pages: int
+def _is_plain_data(value: Any) -> bool:
+ if isinstance(value, dict):
+ return all(isinstance(key, str) and _is_plain_data(item) for key, item
in value.items())
+ if isinstance(value, (list, tuple)):
+ return all(_is_plain_data(item) for item in value)
+ return value is None or isinstance(value, (str, int, float, bool))
+
+
+def build_start_trigger_args(
+ trigger: MSGraphTrigger, next_method: str, timeout: timedelta | None = None
+) -> StartTriggerArgs | None:
+ """
+ Build the arguments which start a task from the triggerer with the given
trigger, when that is possible.
+
+ The trigger is the one the task would defer to on a worker, built from its
arguments before they are
+ rendered, so both paths send the same first request. The triggerer
renders the templated fields of
+ that trigger itself, and it only gets what the serialized Dag holds: it
never resolves an ``XComArg``,
+ and a callable or a file-like object cannot be serialized. None is
returned in those cases, and on
+ Airflow before 3.4, so the task starts on a worker instead. Before 3.4
the XComs of an earlier try
+ are not cleared for a task the scheduler deferred, and the operator would
read them back.
+
+ :param trigger: The trigger which sends the first request.
+ :param next_method: The method of the task which handles the event of the
trigger.
+ :param timeout: How long the trigger may run before the task times out
(default is None).
+ """
+ if not AIRFLOW_V_3_4_PLUS:
+ return None
+
+ trigger_cls, trigger_kwargs = trigger.serialize()
+ trigger_kwargs["api_version"] =
KiotaRequestAdapterHook.resolve_api_version_from_value(
+ trigger_kwargs["api_version"]
+ )
+ # This trigger is the one the triggerer renders, unlike the ones a worker
defers afterwards.
+ trigger_kwargs["rendered"] = False
+
+ if not _is_plain_data(trigger_kwargs):
+ return None
+
+ return StartTriggerArgs(
+ trigger_cls=trigger_cls,
+ trigger_kwargs=trigger_kwargs,
+ next_method=next_method,
+ timeout=timeout,
+ )
+
+
def default_event_handler(event: dict[Any, Any] | None = None, **context) ->
Any:
if event:
if event.get("status") == "failure":
@@ -137,6 +185,10 @@ class MSGraphAsyncOperator(BaseOperator):
are removed from the store once the result is returned. Pages are
stored as they are, so pages
larger than the store's ``max_value_storage_bytes`` need a state store
backend that offloads
large values. Needs Airflow 3.3 or later.
+ :param start_from_trigger: Start the task directly from the triggerer,
without first running it on a
+ worker (default is False). This requires Airflow 3.4 or later, and
arguments which only hold plain
+ values: when an argument is an ``XComArg``, a callable or a file-like
object, or when ``durable`` is
+ set, the task starts on a worker as usual. The templated fields are
then rendered by the triggerer.
"""
# Durable execution implemented directly against task_state_store, without
ResumableJobMixin.
@@ -175,6 +227,7 @@ class MSGraphAsyncOperator(BaseOperator):
event_handler: Callable[[dict[Any, Any] | None, Context], Any] | None
= None,
serializer: type[ResponseSerializer] = ResponseSerializer,
durable: bool = False,
+ start_from_trigger: bool = False,
**kwargs: Any,
):
super().__init__(**kwargs)
@@ -200,6 +253,14 @@ class MSGraphAsyncOperator(BaseOperator):
self.event_handler = event_handler or default_event_handler
self.serializer: ResponseSerializer = serializer()
self.durable = durable
+ # A durable task resumes from the task state store in execute, which a
task the scheduler
+ # deferred never runs.
+ self.start_trigger_args = (
+ build_start_trigger_args(self._build_trigger(),
next_method=self.execute_complete.__name__)
+ if start_from_trigger and not durable
+ else None
+ )
+ self.start_from_trigger = self.start_trigger_args is not None
def execute(self, context: Context) -> Any:
if task_state_store := self.get_task_state_store(context):
@@ -225,24 +286,26 @@ class MSGraphAsyncOperator(BaseOperator):
"task_state_store not available in context, crash recovery is
disabled for this run"
)
- self.defer(
- trigger=MSGraphTrigger(
- url=self.url,
- response_type=self.response_type,
- path_parameters=self.path_parameters,
- url_template=self.url_template,
- method=self.method,
- query_parameters=self.query_parameters,
- headers=self.headers,
- data=self.data,
- conn_id=self.conn_id,
- timeout=self.timeout,
- proxies=self.proxies,
- scopes=self.scopes,
- api_version=self.api_version,
- serializer=type(self.serializer),
- ),
- method_name=self.execute_complete.__name__,
+ self.defer(trigger=self._build_trigger(),
method_name=self.execute_complete.__name__)
+
+ def _build_trigger(self) -> MSGraphTrigger:
+ """Build the trigger which sends the first request, with the arguments
as they are now."""
+ return MSGraphTrigger(
+ url=self.url,
+ response_type=self.response_type,
+ path_parameters=self.path_parameters,
+ url_template=self.url_template,
+ method=self.method,
+ query_parameters=self.query_parameters,
+ headers=self.headers,
+ data=self.data,
+ conn_id=self.conn_id,
+ timeout=self.timeout,
+ proxies=self.proxies,
+ scopes=self.scopes,
+ api_version=self.api_version,
+ serializer=type(self.serializer),
+ rendered=True,
)
def execute_complete(
@@ -532,6 +595,7 @@ class MSGraphAsyncOperator(BaseOperator):
api_version=self.api_version,
serializer=type(self.serializer),
pagination_link=True,
+ rendered=True,
),
method_name=method_name,
kwargs={"query_parameters": query_parameters},
diff --git
a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/sensors/msgraph.py
b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/sensors/msgraph.py
index beada87b7a1..3b0dd5bad4a 100644
---
a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/sensors/msgraph.py
+++
b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/sensors/msgraph.py
@@ -29,7 +29,7 @@ from airflow.providers.common.compat.sdk import (
)
from airflow.providers.common.compat.standard.triggers import TimeDeltaTrigger
from airflow.providers.microsoft.azure.hooks.msgraph import
KiotaRequestAdapterHook
-from airflow.providers.microsoft.azure.operators.msgraph import
execute_callable
+from airflow.providers.microsoft.azure.operators.msgraph import
build_start_trigger_args, execute_callable
from airflow.providers.microsoft.azure.triggers.msgraph import MSGraphTrigger,
ResponseSerializer
if TYPE_CHECKING:
@@ -62,6 +62,10 @@ class MSGraphSensor(BaseSensorOperator):
`KiotaRequestAdapterHook` are bytes, then those will be base64 encoded
into a string.
:param serializer: Class which handles response serialization (default is
ResponseSerializer).
Bytes will be base64 encoded into a string, so it can be stored as an
XCom.
+ :param start_from_trigger: Start the task directly from the triggerer,
without first running it on a
+ worker (default is False). This requires Airflow 3.4 or later, and
arguments which only hold plain
+ values: when an argument is an ``XComArg``, a callable or a file-like
object, the task starts on a
+ worker as usual. The templated fields are then rendered by the
triggerer.
"""
template_fields: Sequence[str] = (
@@ -93,6 +97,7 @@ class MSGraphSensor(BaseSensorOperator):
result_processor: Callable[[Any, Context], Any] = lambda result,
**context: result,
serializer: type[ResponseSerializer] = ResponseSerializer,
retry_delay: timedelta | float = 60,
+ start_from_trigger: bool = False,
**kwargs,
):
super().__init__(retry_delay=retry_delay, **kwargs)
@@ -111,40 +116,62 @@ class MSGraphSensor(BaseSensorOperator):
self.event_processor = event_processor
self.result_processor = result_processor
self.serializer = serializer()
+ # The scheduler defers the task itself here, and it only knows the
timeout it is handed.
+ self.start_trigger_args = (
+ build_start_trigger_args(
+ self._build_trigger(),
+ next_method=self.execute_complete.__name__,
+ timeout=timedelta(seconds=self.timeout),
+ )
+ if start_from_trigger
+ else None
+ )
+ self.start_from_trigger = self.start_trigger_args is not None
def execute(self, context: Context):
self._defer_request(timezone.utcnow().timestamp(), self.timeout)
def retry_execute(self, context: Context, sensor_started_at: float | None
= None, **kwargs) -> Any:
- started_at = self._started_at(sensor_started_at)
+ started_at = self._started_at(context, sensor_started_at)
remaining = max(self._remaining_timeout(started_at),
self.retry_delay.total_seconds())
self._defer_request(started_at, remaining)
def _defer_request(self, started_at: float, remaining: float) -> None:
self.defer(
- trigger=MSGraphTrigger(
- url=self.url,
- response_type=self.response_type,
- path_parameters=self.path_parameters,
- url_template=self.url_template,
- method=self.method,
- query_parameters=self.query_parameters,
- headers=self.headers,
- data=self.data,
- conn_id=self.conn_id,
- proxies=self.proxies,
- scopes=self.scopes,
- api_version=self.api_version,
- serializer=type(self.serializer),
- ),
+ trigger=self._build_trigger(),
method_name=self.execute_complete.__name__,
kwargs={"sensor_started_at": started_at},
timeout=timedelta(seconds=remaining),
)
+ def _build_trigger(self) -> MSGraphTrigger:
+ """Build the trigger which polls once, with the arguments as they are
now."""
+ return MSGraphTrigger(
+ url=self.url,
+ response_type=self.response_type,
+ path_parameters=self.path_parameters,
+ url_template=self.url_template,
+ method=self.method,
+ query_parameters=self.query_parameters,
+ headers=self.headers,
+ data=self.data,
+ conn_id=self.conn_id,
+ proxies=self.proxies,
+ scopes=self.scopes,
+ api_version=self.api_version,
+ serializer=type(self.serializer),
+ rendered=True,
+ )
+
@staticmethod
- def _started_at(sensor_started_at: float | None) -> float:
- return sensor_started_at if sensor_started_at is not None else
timezone.utcnow().timestamp()
+ def _started_at(context: Context, sensor_started_at: float | None) ->
float:
+ if sensor_started_at is not None:
+ return sensor_started_at
+ # The first poll of a task the scheduler deferred brings no
sensor_started_at: the scheduler
+ # set the start of the try when it deferred the task.
+ if start_date := getattr(context.get("ti"), "start_date", None):
+ return start_date.timestamp()
+ return timezone.utcnow().timestamp()
def _remaining_timeout(self, started_at: float) -> float:
return self.timeout - max(timezone.utcnow().timestamp() - started_at,
0)
@@ -198,7 +225,7 @@ class MSGraphSensor(BaseSensorOperator):
return result
- started_at = self._started_at(sensor_started_at)
+ started_at = self._started_at(context, sensor_started_at)
remaining = self._remaining_timeout(started_at)
if remaining <= 0:
raise AirflowSensorTimeout(
diff --git
a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/triggers/msgraph.py
b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/triggers/msgraph.py
index b6c32bb3f1d..f55c436fa90 100644
---
a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/triggers/msgraph.py
+++
b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/triggers/msgraph.py
@@ -40,9 +40,12 @@ from airflow.triggers.base import BaseTrigger, TriggerEvent
if TYPE_CHECKING:
from io import BytesIO
+ import jinja2
from kiota_abstractions.request_adapter import RequestAdapter
from msgraph_core import APIVersion
+ from airflow.sdk import Context
+
class ResponseSerializer:
"""ResponseSerializer serializes the response as a string."""
@@ -98,6 +101,10 @@ class MSGraphTrigger(BaseTrigger):
:param pagination_link: Whether `url` was taken from a pagination link of
a previous response
(default is False). When True, its host is verified against the
configured Microsoft Graph
endpoint before the request is made.
+ :param rendered: Whether the templated fields already hold their rendered
values (default is False).
+ A task which starts from the triggerer gets every one of its triggers
rendered by the triggerer,
+ also the ones a worker deferred with values it rendered itself or took
from a response, such as
+ a pagination link. When True, the triggerer leaves the fields as they
are.
"""
def __init__(
@@ -117,6 +124,7 @@ class MSGraphTrigger(BaseTrigger):
api_version: APIVersion | str | None = None,
serializer: type[ResponseSerializer] = ResponseSerializer,
pagination_link: bool = False,
+ rendered: bool = False,
):
super().__init__()
self.conn_id = conn_id
@@ -134,6 +142,7 @@ class MSGraphTrigger(BaseTrigger):
self.data = data
self.serializer: ResponseSerializer = self.resolve_type(serializer,
default=ResponseSerializer)()
self.pagination_link = pagination_link
+ self.rendered = rendered
@classmethod
def resolve_type(cls, value: str | type, default) -> type:
@@ -163,9 +172,14 @@ class MSGraphTrigger(BaseTrigger):
"data": self.data,
"response_type": self.response_type,
"pagination_link": self.pagination_link,
+ "rendered": self.rendered,
},
)
+ def render_template_fields(self, context: Context, jinja_env:
jinja2.Environment | None = None) -> None:
+ if not self.rendered:
+ super().render_template_fields(context=context,
jinja_env=jinja_env)
+
def get_conn(self) -> RequestAdapter:
"""
Initiate a new RequestAdapter connection.
diff --git
a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/version_compat.py
b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/version_compat.py
index 8b5cdcf5cc4..67f7e4c3e50 100644
---
a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/version_compat.py
+++
b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/version_compat.py
@@ -35,9 +35,11 @@ def get_base_airflow_version_tuple() -> tuple[int, int, int]:
AIRFLOW_V_3_0_PLUS = get_base_airflow_version_tuple() >= (3, 0, 0)
AIRFLOW_V_3_1_PLUS = get_base_airflow_version_tuple() >= (3, 1, 0)
AIRFLOW_V_3_3_PLUS = get_base_airflow_version_tuple() >= (3, 3, 0)
+AIRFLOW_V_3_4_PLUS = get_base_airflow_version_tuple() >= (3, 4, 0)
__all__ = [
"AIRFLOW_V_3_0_PLUS",
"AIRFLOW_V_3_1_PLUS",
"AIRFLOW_V_3_3_PLUS",
+ "AIRFLOW_V_3_4_PLUS",
]
diff --git
a/providers/microsoft/azure/tests/unit/microsoft/azure/operators/test_msgraph.py
b/providers/microsoft/azure/tests/unit/microsoft/azure/operators/test_msgraph.py
index 3d6ace80136..d44c63a97c4 100644
---
a/providers/microsoft/azure/tests/unit/microsoft/azure/operators/test_msgraph.py
+++
b/providers/microsoft/azure/tests/unit/microsoft/azure/operators/test_msgraph.py
@@ -20,14 +20,17 @@ import json
import locale
import warnings
from base64 import b64encode
+from io import BytesIO
from os.path import dirname
from typing import Any
from unittest import mock
import pytest
+from msgraph_core import APIVersion
from airflow.exceptions import AirflowProviderDeprecationWarning
-from airflow.providers.common.compat.sdk import AirflowException, Context,
TaskDeferred, timezone
+from airflow.models.trigger import Trigger
+from airflow.providers.common.compat.sdk import DAG, AirflowException,
Context, TaskDeferred, timezone
from airflow.providers.microsoft.azure.exceptions import
MSGraphPageNotInTaskStateStoreError
from airflow.providers.microsoft.azure.operators.msgraph import (
_COMPLETED_KEY,
@@ -36,12 +39,14 @@ from airflow.providers.microsoft.azure.operators.msgraph
import (
MSGraphAsyncOperator,
execute_callable,
)
+from airflow.providers.microsoft.azure.triggers.msgraph import MSGraphTrigger
from airflow.triggers.base import TriggerEvent
+from airflow.utils.state import TaskInstanceState
from tests_common.test_utils.file_loading import load_file_from_resources,
load_json_from_resources
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_3_3_PLUS
+from tests_common.test_utils.operators.run_deferrable import execute_operator,
run_trigger
+from tests_common.test_utils.version_compat import AIRFLOW_V_3_3_PLUS,
AIRFLOW_V_3_4_PLUS
from unit.microsoft.azure.test_utils import (
mock_json_response,
mock_response,
@@ -783,3 +788,279 @@ class TestDurableMSGraphAsyncOperator:
mock_log.return_value.warning.assert_called_once_with(
"task_state_store not available in context, crash recovery is
disabled for this run"
)
+
+
[email protected](not AIRFLOW_V_3_4_PLUS, reason="A task starts from the
triggerer from Airflow 3.4")
+class TestMSGraphAsyncOperatorStartFromTrigger:
+ def test_start_from_trigger_builds_the_trigger_arguments(self):
+ operator = MSGraphAsyncOperator(
+ task_id="user_messages",
+ conn_id="{{ params.conn_id }}",
+ url="users/{user_id}/messages",
+ response_type="bytes",
+ path_parameters={"user_id": "{{ params.user_id }}"},
+ url_template="{+baseurl}/users/{user_id}/messages{?%24top}",
+ method="POST",
+ query_parameters={"$top": 12},
+ headers={"ConsistencyLevel": "eventual"},
+ data={"subject": "{{ ds }}"},
+ timeout=30,
+ proxies={"https": "http://proxy:3128"},
+ scopes=["https://graph.microsoft.com/.default";],
+ api_version=APIVersion.beta,
+ start_from_trigger=True,
+ )
+
+ assert operator.start_from_trigger is True
+ assert operator.start_trigger_args.trigger_cls == (
+ "airflow.providers.microsoft.azure.triggers.msgraph.MSGraphTrigger"
+ )
+ assert operator.start_trigger_args.next_method == "execute_complete"
+ assert operator.start_trigger_args.next_kwargs is None
+ assert operator.start_trigger_args.timeout is None
+ assert operator.start_trigger_args.trigger_kwargs == {
+ "url": "users/{user_id}/messages",
+ "response_type": "bytes",
+ "path_parameters": {"user_id": "{{ params.user_id }}"},
+ "url_template": "{+baseurl}/users/{user_id}/messages{?%24top}",
+ "method": "POST",
+ "query_parameters": {"$top": 12},
+ "headers": {"ConsistencyLevel": "eventual"},
+ "data": {"subject": "{{ ds }}"},
+ "conn_id": "{{ params.conn_id }}",
+ "timeout": 30,
+ "proxies": {"https": "http://proxy:3128"},
+ "scopes": ["https://graph.microsoft.com/.default";],
+ "api_version": "beta",
+ "serializer":
"airflow.providers.microsoft.azure.triggers.msgraph.ResponseSerializer",
+ "pagination_link": False,
+ "rendered": False,
+ }
+ # The enum itself would not survive the serialized Dag, its value does.
+ assert type(operator.start_trigger_args.trigger_kwargs["api_version"])
is str
+
+ def test_start_from_trigger_defers_the_same_trigger_as_execute(self):
+ operator = MSGraphAsyncOperator(
+ task_id="user_messages",
+ conn_id="msgraph_api",
+ url="users/{user_id}/messages",
+ path_parameters={"user_id": "{{ params.user_id }}"},
+ query_parameters={"$top": 12},
+ api_version=APIVersion.beta,
+ start_from_trigger=True,
+ )
+
+ with pytest.raises(TaskDeferred) as deferred:
+ operator.execute(context=Context())
+
+ trigger = MSGraphTrigger(**operator.start_trigger_args.trigger_kwargs)
+
+ trigger_cls, trigger_kwargs = deferred.value.trigger.serialize()
+ # The worker defers its trigger with rendered values, the triggerer
renders the start trigger itself.
+ assert trigger_kwargs["rendered"] is True
+ assert trigger.serialize() == (trigger_cls, {**trigger_kwargs,
"rendered": False})
+ assert operator.start_trigger_args.next_method ==
deferred.value.method_name
+
+ @pytest.mark.parametrize(
+ "arguments",
+ [
+ pytest.param(lambda upstream: {"url": upstream.output},
id="xcom-arg"),
+ pytest.param(
+ lambda upstream: {"path_parameters": {"user_id":
upstream.output}}, id="nested-xcom-arg"
+ ),
+ pytest.param(
+ lambda upstream: {"query_parameters": {"$select":
[upstream.output]}},
+ id="xcom-arg-in-list",
+ ),
+ pytest.param(
+ lambda upstream: {"path_parameters": lambda context,
jinja_env: {"user_id": "me"}},
+ id="callable",
+ ),
+ pytest.param(lambda upstream: {"data": BytesIO(b"content")},
id="file-like-object"),
+ ],
+ )
+ def test_start_from_trigger_falls_back_to_the_worker(self, arguments):
+ with DAG(dag_id="msgraph_start_from_trigger", schedule=None):
+ upstream = MSGraphAsyncOperator(task_id="upstream",
conn_id="msgraph_api", url="users")
+ operator = MSGraphAsyncOperator(
+ **{
+ "task_id": "user_messages",
+ "conn_id": "msgraph_api",
+ "url": "users/{user_id}/messages",
+ "start_from_trigger": True,
+ **arguments(upstream),
+ }
+ )
+
+ assert operator.start_from_trigger is False
+ assert operator.start_trigger_args is None
+
+
@mock.patch("airflow.providers.microsoft.azure.operators.msgraph.AIRFLOW_V_3_4_PLUS",
False)
+ def test_start_from_trigger_needs_airflow_3_4(self):
+ # Before 3.4 the XComs of an earlier try are not cleared for a task
the scheduler deferred.
+ operator = MSGraphAsyncOperator(
+ task_id="users", conn_id="msgraph_api", url="users",
start_from_trigger=True
+ )
+
+ assert operator.start_from_trigger is False
+ assert operator.start_trigger_args is None
+
+ def test_durable_task_starts_on_a_worker(self):
+ operator = MSGraphAsyncOperator(
+ task_id="users", conn_id="msgraph_api", url="users", durable=True,
start_from_trigger=True
+ )
+
+ assert operator.start_from_trigger is False
+ assert operator.start_trigger_args is None
+
+ def test_triggers_deferred_by_the_worker_are_not_rendered_again(self):
+ operator = MSGraphAsyncOperator(
+ task_id="users", conn_id="msgraph_api", url="users",
start_from_trigger=True
+ )
+
+ with pytest.raises(TaskDeferred) as first_page:
+ operator.execute(context=Context())
+ with pytest.raises(TaskDeferred) as next_page:
+ operator.defer_to_next_link(
+ url="https://graph.microsoft.com/v1.0/users?$skiptoken={{ x
}}",
+ query_parameters=None,
+ method_name=operator.execute_complete.__name__,
+ )
+
+ assert first_page.value.trigger.rendered is True
+ assert next_page.value.trigger.rendered is True
+
+ def test_start_trigger_args_are_not_shared_between_tasks(self):
+ users = MSGraphAsyncOperator(
+ task_id="users", conn_id="msgraph_api", url="users",
start_from_trigger=True
+ )
+ groups = MSGraphAsyncOperator(
+ task_id="groups", conn_id="msgraph_api", url="groups",
start_from_trigger=True
+ )
+ sites = MSGraphAsyncOperator(task_id="sites", conn_id="msgraph_api",
url="sites")
+
+ assert users.start_trigger_args is not groups.start_trigger_args
+ assert users.start_trigger_args.trigger_kwargs["url"] == "users"
+ assert groups.start_trigger_args.trigger_kwargs["url"] == "groups"
+ assert sites.start_from_trigger is False
+ assert sites.start_trigger_args is None
+ assert MSGraphAsyncOperator.start_from_trigger is False
+ assert MSGraphAsyncOperator.start_trigger_args is None
+
+ def test_mapped_task_has_no_start_trigger_args(self):
+ with DAG(dag_id="msgraph_start_from_trigger", schedule=None):
+ mapped = MSGraphAsyncOperator.partial(
+ task_id="users", conn_id="msgraph_api", start_from_trigger=True
+ ).expand(url=["users", "groups"])
+
+ # The scheduler cannot expand the arguments of a mapped task, and
without them it leaves
+ # the task to a worker.
+ assert mapped.start_trigger_args is None
+
+ @pytest.mark.db_test
+ @pytest.mark.need_serialized_dag
+ def test_scheduler_defers_the_task_to_the_triggerer(self, dag_maker,
session):
+ with dag_maker(session=session):
+ users = MSGraphAsyncOperator(
+ task_id="users",
+ conn_id="msgraph_api",
+ url="users/{{ params.user_id }}",
+ query_parameters={"$top": 12},
+ start_from_trigger=True,
+ )
+ MSGraphAsyncOperator(
+ task_id="messages", conn_id="msgraph_api", url=users.output,
start_from_trigger=True
+ )
+
+ dag_run = dag_maker.create_dagrun()
+ task_instances = {}
+ for task_id in ("users", "messages"):
+ task_instances[task_id] = dag_run.get_task_instance(task_id,
session=session)
+ task_instances[task_id].task = dag_run.dag.get_task(task_id)
+
+ dag_run.schedule_tis(task_instances.values(), session=session)
+
+ assert task_instances["users"].state == TaskInstanceState.DEFERRED
+ assert task_instances["users"].next_method == "execute_complete"
+ assert task_instances["users"].trigger_timeout is None
+ trigger = session.get(Trigger, task_instances["users"].trigger_id)
+ assert trigger.classpath ==
"airflow.providers.microsoft.azure.triggers.msgraph.MSGraphTrigger"
+ assert trigger.kwargs == users.start_trigger_args.trigger_kwargs
+ # The task which takes its url from an XCom still starts on a worker.
+ session.refresh(task_instances["messages"])
+ assert task_instances["messages"].state == TaskInstanceState.SCHEDULED
+ assert task_instances["messages"].trigger_id is None
+
+ @pytest.mark.db_test
+ def test_triggerer_renders_the_templated_fields(self,
create_task_instance):
+ users = load_json_from_resources(dirname(__file__), "..", "resources",
"users.json")
+ users.pop("@odata.nextLink")
+ operator = MSGraphAsyncOperator(
+ task_id="user_messages",
+ conn_id="{{ conn_id }}",
+ url="users/{user_id}/messages",
+ path_parameters={"user_id": "{{ user_id }}"},
+ query_parameters={"$top": "{{ top }}"},
+ headers={"ConsistencyLevel": "{{ consistency }}"},
+ start_from_trigger=True,
+ )
+ task_instance = create_task_instance(
+ task=operator, start_from_trigger=True,
start_trigger_args=operator.start_trigger_args
+ )
+ trigger = MSGraphTrigger(**operator.start_trigger_args.trigger_kwargs)
+
+ trigger.task_instance = task_instance
+
+ assert trigger.template_fields == MSGraphAsyncOperator.template_fields
+
+ trigger.render_template_fields(
+ context={"conn_id": "msgraph_api", "user_id": "me", "top": 12,
"consistency": "eventual"}
+ )
+
+ assert trigger.conn_id == "msgraph_api"
+ assert trigger.path_parameters == {"user_id": "me"}
+ assert trigger.query_parameters == {"$top": "12"}
+ assert trigger.headers == {"ConsistencyLevel": "eventual"}
+
+ with patch_hook_and_request_adapter(mock_json_response(200, users)) as
(
+ *_,
+ mock_get_http_response,
+ ):
+ events = run_trigger(trigger)
+
+ request = mock_get_http_response.call_args.args[0]
+
+ assert request.url == "users/me/messages?%24top=12"
+ assert request.headers.try_get("ConsistencyLevel") == {"eventual"}
+ assert events[0].payload["status"] == "success"
+ assert events[0].payload["response"] == json.dumps(users)
+
+ def test_execute_when_started_from_the_trigger(self):
+ users = load_json_from_resources(dirname(__file__), "..", "resources",
"users.json")
+ next_users = load_json_from_resources(dirname(__file__), "..",
"resources", "next_users.json")
+ response = mock_json_response(200, users, next_users)
+
+ with (
+ patch_hook_and_request_adapter(response) as (*_,
mock_get_http_response),
+ mock.patch.object(MSGraphAsyncOperator, "execute", autospec=True)
as mock_execute,
+ ):
+ operator = MSGraphAsyncOperator(
+ task_id="users",
+ conn_id="msgraph_api",
+ url="{{ ti.task_id }}",
+ result_processor=lambda result, **context: result.get("value"),
+ start_from_trigger=True,
+ )
+
+ results, events = execute_operator(operator)
+
+ # The first page is requested by the trigger the scheduler created,
with the url it rendered
+ # itself, and the worker only takes over from there to follow the
pagination.
+ mock_execute.assert_not_called()
+ assert mock_get_http_response.call_args_list[0].args[0].url == "users"
+ assert mock_get_http_response.call_count == 2
+ assert results == users.get("value") + next_users.get("value")
+ assert [event.payload["response"] for event in events] == [
+ json.dumps(users),
+ json.dumps(next_users),
+ ]
diff --git
a/providers/microsoft/azure/tests/unit/microsoft/azure/sensors/test_msgraph.py
b/providers/microsoft/azure/tests/unit/microsoft/azure/sensors/test_msgraph.py
index 2dd951c1e99..1c469ba4d8d 100644
---
a/providers/microsoft/azure/tests/unit/microsoft/azure/sensors/test_msgraph.py
+++
b/providers/microsoft/azure/tests/unit/microsoft/azure/sensors/test_msgraph.py
@@ -19,21 +19,32 @@ from __future__ import annotations
import json
from datetime import UTC, datetime, timedelta
from os.path import dirname
+from types import SimpleNamespace
from unittest.mock import patch
import pytest
import time_machine
from airflow.exceptions import AirflowProviderDeprecationWarning
-from airflow.providers.common.compat.sdk import AirflowSensorTimeout,
AirflowSkipException, TaskDeferred
+from airflow.models.trigger import Trigger
+from airflow.providers.common.compat.sdk import (
+ DAG,
+ AirflowSensorTimeout,
+ AirflowSkipException,
+ Context,
+ TaskDeferred,
+ timezone,
+)
from airflow.providers.common.compat.standard.triggers import TimeDeltaTrigger
+from airflow.providers.microsoft.azure.operators.msgraph import
MSGraphAsyncOperator
from airflow.providers.microsoft.azure.sensors.msgraph import MSGraphSensor
from airflow.providers.microsoft.azure.triggers.msgraph import MSGraphTrigger
from airflow.triggers.base import TriggerEvent
+from airflow.utils.state import TaskInstanceState
from tests_common.test_utils.file_loading import load_json_from_resources
from tests_common.test_utils.operators.run_deferrable import execute_operator
-from tests_common.test_utils.version_compat import AIRFLOW_V_3_1_PLUS
+from tests_common.test_utils.version_compat import AIRFLOW_V_3_1_PLUS,
AIRFLOW_V_3_4_PLUS
from unit.microsoft.azure.test_utils import mock_json_response,
patch_hook_and_request_adapter
T0 = datetime(2026, 1, 1, tzinfo=UTC)
@@ -268,6 +279,20 @@ class TestMSGraphSensor:
"execute_complete", {"event": NOT_DONE, "sensor_started_at":
T0.timestamp()}, {}
)
+ def
test_first_poll_from_the_triggerer_counts_from_the_start_of_the_try(self):
+ sensor = _sensor()
+ task_instance = SimpleNamespace(start_date=T0)
+
+ with _at(50), pytest.raises(TaskDeferred) as exc:
+ sensor.execute_complete(context={"ti": task_instance},
event=NOT_DONE)
+
+ assert exc.value.kwargs == {"sensor_started_at": T0.timestamp()}
+ assert exc.value.trigger.moment == T0 + timedelta(seconds=80)
+
+ # Counted from the start of the try, the 100 seconds are over, not
from the first event.
+ with _at(150), pytest.raises(AirflowSensorTimeout):
+ sensor.execute_complete(context={"ti": task_instance},
event=NOT_DONE)
+
def test_missing_started_at_falls_back_to_now(self):
sensor = _sensor()
@@ -292,3 +317,175 @@ class TestMSGraphSensor:
)
assert exc.value.timeout == timedelta(seconds=100)
+
+
[email protected](not AIRFLOW_V_3_4_PLUS, reason="A task starts from the
triggerer from Airflow 3.4")
+class TestMSGraphSensorStartFromTrigger:
+ def test_start_from_trigger_builds_the_trigger_arguments(self):
+ sensor = MSGraphSensor(
+ task_id="check_workspaces_status",
+ conn_id="powerbi",
+ url="myorg/admin/workspaces/scanStatus/{scanId}",
+ path_parameters={"scanId": "{{ params.scan_id }}"},
+ timeout=10,
+ start_from_trigger=True,
+ )
+
+ assert sensor.start_from_trigger is True
+ assert sensor.start_trigger_args.trigger_cls == (
+ "airflow.providers.microsoft.azure.triggers.msgraph.MSGraphTrigger"
+ )
+ assert sensor.start_trigger_args.next_method == "execute_complete"
+ assert sensor.start_trigger_args.timeout == timedelta(seconds=10)
+ assert sensor.start_trigger_args.trigger_kwargs == {
+ "url": "myorg/admin/workspaces/scanStatus/{scanId}",
+ "response_type": None,
+ "path_parameters": {"scanId": "{{ params.scan_id }}"},
+ "url_template": None,
+ "method": "GET",
+ "query_parameters": None,
+ "headers": None,
+ "data": None,
+ "conn_id": "powerbi",
+ "timeout": None,
+ "proxies": None,
+ "scopes": None,
+ "api_version": None,
+ "serializer":
"airflow.providers.microsoft.azure.triggers.msgraph.ResponseSerializer",
+ "pagination_link": False,
+ "rendered": False,
+ }
+
+ def test_start_from_trigger_defers_the_same_trigger_as_execute(self):
+ sensor = MSGraphSensor(
+ task_id="check_workspaces_status",
+ conn_id="powerbi",
+ url="myorg/admin/workspaces/scanStatus/{scanId}",
+ path_parameters={"scanId": "{{ params.scan_id }}"},
+ timeout=10,
+ start_from_trigger=True,
+ )
+
+ with pytest.raises(TaskDeferred) as deferred:
+ sensor.execute(context=Context())
+
+ trigger = MSGraphTrigger(**sensor.start_trigger_args.trigger_kwargs)
+
+ trigger_cls, trigger_kwargs = deferred.value.trigger.serialize()
+ # The worker defers its trigger with rendered values, the triggerer
renders the start trigger itself.
+ assert trigger_kwargs["rendered"] is True
+ assert trigger.serialize() == (trigger_cls, {**trigger_kwargs,
"rendered": False})
+ assert sensor.start_trigger_args.next_method ==
deferred.value.method_name
+
+ def test_start_from_trigger_falls_back_to_the_worker(self):
+ with DAG(dag_id="msgraph_start_from_trigger", schedule=None):
+ scan = MSGraphAsyncOperator(task_id="scan", conn_id="powerbi",
url="myorg/admin/workspaces")
+ sensor = MSGraphSensor(
+ task_id="check_workspaces_status",
+ conn_id="powerbi",
+ url="myorg/admin/workspaces/scanStatus/{scanId}",
+ path_parameters={"scanId": scan.output},
+ start_from_trigger=True,
+ )
+
+ assert sensor.start_from_trigger is False
+ assert sensor.start_trigger_args is None
+
+ def test_start_trigger_args_are_not_shared_between_tasks(self):
+ first = MSGraphSensor(
+ task_id="first", conn_id="powerbi", url="first", timeout=10,
start_from_trigger=True
+ )
+ second = MSGraphSensor(
+ task_id="second", conn_id="powerbi", url="second", timeout=20,
start_from_trigger=True
+ )
+ third = MSGraphSensor(task_id="third", conn_id="powerbi", url="third")
+
+ assert first.start_trigger_args is not second.start_trigger_args
+ assert first.start_trigger_args.trigger_kwargs["url"] == "first"
+ assert first.start_trigger_args.timeout == timedelta(seconds=10)
+ assert second.start_trigger_args.trigger_kwargs["url"] == "second"
+ assert second.start_trigger_args.timeout == timedelta(seconds=20)
+ assert third.start_from_trigger is False
+ assert third.start_trigger_args is None
+ assert MSGraphSensor.start_from_trigger is False
+ assert MSGraphSensor.start_trigger_args is None
+
+ def test_mapped_task_has_no_start_trigger_args(self):
+ with DAG(dag_id="msgraph_start_from_trigger", schedule=None):
+ mapped = MSGraphSensor.partial(
+ task_id="check_workspaces_status",
+ conn_id="powerbi",
+ url="myorg/admin/workspaces/scanStatus/{scanId}",
+ start_from_trigger=True,
+ ).expand(path_parameters=[{"scanId": "first"}, {"scanId":
"second"}])
+
+ # The scheduler cannot expand the arguments of a mapped task, and
without them it leaves
+ # the task to a worker.
+ assert mapped.start_trigger_args is None
+
+ @pytest.mark.db_test
+ @pytest.mark.need_serialized_dag
+ def test_scheduler_defers_the_task_with_the_sensor_timeout(self,
dag_maker, session, time_machine):
+ now = timezone.datetime(2026, 10, 5, 12)
+ time_machine.move_to(now, tick=False)
+
+ with dag_maker(session=session):
+ sensor = MSGraphSensor(
+ task_id="check_workspaces_status",
+ conn_id="powerbi",
+ url="myorg/admin/workspaces/scanStatus/{scanId}",
+ path_parameters={"scanId": "{{ params.scan_id }}"},
+ timeout=10,
+ start_from_trigger=True,
+ )
+
+ dag_run = dag_maker.create_dagrun()
+ task_instance = dag_run.get_task_instance("check_workspaces_status",
session=session)
+ task_instance.task = dag_run.dag.get_task("check_workspaces_status")
+
+ dag_run.schedule_tis((task_instance,), session=session)
+
+ assert task_instance.state == TaskInstanceState.DEFERRED
+ assert task_instance.next_method == "execute_complete"
+ assert task_instance.trigger_timeout == now + timedelta(seconds=10)
+ trigger = session.get(Trigger, task_instance.trigger_id)
+ assert trigger.classpath ==
"airflow.providers.microsoft.azure.triggers.msgraph.MSGraphTrigger"
+ assert trigger.kwargs == sensor.start_trigger_args.trigger_kwargs
+
+ def test_execute_when_started_from_the_trigger(self):
+ status = load_json_from_resources(dirname(__file__), "..",
"resources", "status.json")
+ response = mock_json_response(200, *status)
+
+ with (
+ patch_hook_and_request_adapter(response) as (*_,
mock_get_http_response),
+ patch.object(MSGraphSensor, "execute", autospec=True) as
mock_execute,
+ patch.object(
+ MSGraphSensor, "retry_execute", autospec=True,
side_effect=MSGraphSensor.retry_execute
+ ) as mock_retry_execute,
+ ):
+ sensor = MSGraphSensor(
+ task_id="check_workspaces_status",
+ conn_id="powerbi",
+ url="myorg/admin/workspaces/scanStatus/{scanId}",
+ path_parameters={"scanId": "{{ ti.task_id }}"},
+ result_processor=lambda result, **context: result["id"],
+ retry_delay=1,
+ timeout=5,
+ start_from_trigger=True,
+ )
+
+ results, events = execute_operator(sensor)
+
+ # The first poll comes from the trigger the scheduler created, with
the path parameters it
+ # rendered itself, so execute never runs and retry_execute does the
poll after the retry delay.
+ mock_execute.assert_not_called()
+ mock_retry_execute.assert_called_once()
+ assert [call.args[0].url for call in
mock_get_http_response.call_args_list] == [
+ "myorg/admin/workspaces/scanStatus/check_workspaces_status",
+ "myorg/admin/workspaces/scanStatus/check_workspaces_status",
+ ]
+ assert results == "0a1b1bf3-37de-48f7-9863-ed4cda97a9ef"
+ assert len(events) == 3
+ assert events[0].payload["response"] == json.dumps(status[0])
+ assert isinstance(events[1].payload, datetime)
+ assert events[2].payload["response"] == json.dumps(status[1])
diff --git
a/providers/microsoft/azure/tests/unit/microsoft/azure/triggers/test_msgraph.py
b/providers/microsoft/azure/tests/unit/microsoft/azure/triggers/test_msgraph.py
index 7ea42c7df49..02d1ee52176 100644
---
a/providers/microsoft/azure/tests/unit/microsoft/azure/triggers/test_msgraph.py
+++
b/providers/microsoft/azure/tests/unit/microsoft/azure/triggers/test_msgraph.py
@@ -38,6 +38,7 @@ from airflow.triggers.base import TriggerEvent
from tests_common.test_utils.file_loading import load_file_from_resources,
load_json_from_resources
from tests_common.test_utils.operators.run_deferrable import run_trigger
+from tests_common.test_utils.version_compat import AIRFLOW_V_3_3_PLUS
from unit.microsoft.azure.test_utils import (
mock_json_response,
mock_response,
@@ -149,6 +150,7 @@ class TestMSGraphTrigger:
"api_version": APIVersion.v1.value,
"serializer":
f"{ResponseSerializer.__module__}.{ResponseSerializer.__name__}",
"pagination_link": False,
+ "rendered": False,
}
def test_get_conn(self):
@@ -165,6 +167,18 @@ class TestMSGraphTrigger:
assert actual.base_url == "https://graph.microsoft.com/v1.0/";
[email protected](not AIRFLOW_V_3_3_PLUS, reason="The triggerer renders
templated fields from Airflow 3.3")
+class TestMSGraphTriggerRendering:
+ @pytest.mark.parametrize(("rendered", "expected"), [(False, "users/me"),
(True, "users/{{ user_id }}")])
+ def test_render_template_fields(self, rendered, expected):
+ trigger = MSGraphTrigger(url="users/{{ user_id }}", rendered=rendered)
+ trigger.template_fields = ("url",)
+
+ trigger.render_template_fields(context={"user_id": "me"})
+
+ assert trigger.url == expected
+
+
class TestResponseSerializer:
def test_serialize_when_bytes_then_base64_encoded(self):
response = load_file_from_resources(
