Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Added test-case for callable values in path and query parameters of MSGraphAsyncOperator #43799

Merged
Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
31 commits
Select commit Hold shift + click to select a range
a8d6afb
refactor: Allow callable values in path en query parameters to be spe…
davidblain-infrabel Nov 7, 2024
92d3a72
Merge branch 'main' into feature/allow-callable-values-in-query-path-…
dabla Nov 7, 2024
eb5fdb1
refactor: Re-organized imports
davidblain-infrabel Nov 7, 2024
f48b8da
refactor: Reformatted test_evaluate_parameters
davidblain-infrabel Nov 7, 2024
dadb52a
refactor: parameters of evaluate_parameters is optional
davidblain-infrabel Nov 7, 2024
c36680a
refactor: Reformatted test_evaluate_parameters
davidblain-infrabel Nov 8, 2024
b038c44
refactor: Ignore mypy error
davidblain-infrabel Nov 8, 2024
78a6dfd
refactor: Use list comprehension to calculate total length
davidblain-infrabel Nov 8, 2024
b846684
Merge branch 'main' into feature/allow-callable-values-in-query-path-…
dabla Nov 8, 2024
7bb21b6
Merge branch 'main' into feature/allow-callable-values-in-query-path-…
dabla Nov 8, 2024
180ded5
Merge branch 'main' into feature/allow-callable-values-in-query-path-…
dabla Nov 10, 2024
70b87cb
Merge branch 'main' into feature/allow-callable-values-in-query-path-…
dabla Nov 12, 2024
3ce3370
refactor: Removed evaluate_parameters method from KiotaRequestAdapter…
davidblain-infrabel Nov 12, 2024
c01f152
Merge branch 'main' into feature/allow-callable-values-in-query-path-…
dabla Nov 12, 2024
44410ea
refactor: Reorganized imports in KiotaRequestAdapterHook
davidblain-infrabel Nov 12, 2024
c21adc3
refactor: Reorganized imports in TestKiotaRequestAdapterHook
davidblain-infrabel Nov 12, 2024
3d422dd
refactor: Removed obsolete jinja2 import from MSGraphAsyncOperator an…
davidblain-infrabel Nov 12, 2024
e5f8411
refactor: Removed Callable types from path_parameters and query_param…
davidblain-infrabel Nov 12, 2024
e70c852
refactor: Changed type for xcom_pull method of MockedTaskInstance
davidblain-infrabel Nov 12, 2024
f6bc943
Merge branch 'main' into feature/allow-callable-values-in-query-path-…
dabla Nov 12, 2024
86ac427
refactor: Removed duplicate run_id in xcom_pull method of MockedTaskI…
davidblain-infrabel Nov 12, 2024
61bf145
Merge branch 'main' into feature/allow-callable-values-in-query-path-…
dabla Nov 12, 2024
483d582
refactor: Only test operator and sensor with lambda parameter when Ai…
davidblain-infrabel Nov 13, 2024
efa09db
Merge branch 'main' into feature/allow-callable-values-in-query-path-…
dabla Nov 13, 2024
beda59a
refactor: Reorganized imports TestMSGraphSensor
davidblain-infrabel Nov 13, 2024
75442fd
refactor: Reorganized imports TestMSGraphSensor
davidblain-infrabel Nov 13, 2024
8af1ca1
Merge branch 'main' into feature/allow-callable-values-in-query-path-…
dabla Nov 13, 2024
1d9e24f
refactor: Reorganized imports TestMSGraphAsyncOperator
davidblain-infrabel Nov 13, 2024
4b65d82
Merge branch 'main' into feature/allow-callable-values-in-query-path-…
dabla Nov 15, 2024
931fb3d
Merge branch 'main' into feature/allow-callable-values-in-query-path-…
dabla Nov 15, 2024
ec71582
Merge branch 'main' into feature/allow-callable-values-in-query-path-…
dabla Nov 18, 2024
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -280,7 +280,7 @@ def paginate(
if top and odata_count:
if len(response.get("value", [])) == top and context:
results = operator.pull_xcom(context=context)
skip = sum(map(lambda result: len(result["value"]), results)) + top if results else top
skip = sum([len(result["value"]) for result in results]) + top if results else top # type: ignore
query_parameters["$skip"] = skip
return operator.url, query_parameters
return response.get("@odata.nextLink"), operator.query_parameters
Expand Down
1 change: 1 addition & 0 deletions providers/tests/microsoft/azure/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -68,6 +68,7 @@ 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
Expand Down
32 changes: 31 additions & 1 deletion providers/tests/microsoft/azure/operators/test_msgraph.py
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,7 @@
mock_json_response,
mock_response,
)
from tests_common.test_utils.compat import AIRFLOW_V_2_10_PLUS

if TYPE_CHECKING:
from airflow.utils.context import Context
Expand Down Expand Up @@ -143,11 +144,40 @@ def test_execute_when_response_is_bytes(self):
task_id="drive_item_content",
conn_id="msgraph_api",
response_type="bytes",
url=f"/drives/{drive_id}/root/content",
url="/drives/{drive_id}/root/content",
path_parameters={"drive_id": drive_id},
)

results, events = self.execute_operator(operator)

assert operator.path_parameters == {"drive_id": drive_id}
assert results == base64_encoded_content
assert len(events) == 1
assert isinstance(events[0], TriggerEvent)
assert events[0].payload["status"] == "success"
assert events[0].payload["type"] == "builtins.bytes"
assert events[0].payload["response"] == base64_encoded_content

@pytest.mark.db_test
@pytest.mark.skipif(not AIRFLOW_V_2_10_PLUS, reason="Lambda parameters works in Airflow >= 2.10.0")
def test_execute_with_lambda_parameter_when_response_is_bytes(self):
content = load_file("resources", "dummy.pdf", mode="rb", encoding=None)
base64_encoded_content = b64encode(content).decode(locale.getpreferredencoding())
drive_id = "82f9d24d-6891-4790-8b6d-f1b2a1d0ca22"
response = mock_response(200, content)

with self.patch_hook_and_request_adapter(response):
operator = MSGraphAsyncOperator(
task_id="drive_item_content",
conn_id="msgraph_api",
response_type="bytes",
url="/drives/{drive_id}/root/content",
path_parameters=lambda context, jinja_env: {"drive_id": drive_id},
)

results, events = self.execute_operator(operator)

assert operator.path_parameters == {"drive_id": drive_id}
assert results == base64_encoded_content
assert len(events) == 1
assert isinstance(events[0], TriggerEvent)
Expand Down
30 changes: 30 additions & 0 deletions providers/tests/microsoft/azure/sensors/test_msgraph.py
Original file line number Diff line number Diff line change
Expand Up @@ -18,11 +18,14 @@

import json

import pytest

from airflow.providers.microsoft.azure.sensors.msgraph import MSGraphSensor
from airflow.triggers.base import TriggerEvent

from providers.tests.microsoft.azure.base import Base
from providers.tests.microsoft.conftest import load_json, mock_json_response
from tests_common.test_utils.compat import AIRFLOW_V_2_10_PLUS


class TestMSGraphSensor(Base):
Expand All @@ -42,6 +45,33 @@ def test_execute(self):

results, events = self.execute_operator(sensor)

assert sensor.path_parameters == {"scanId": "0a1b1bf3-37de-48f7-9863-ed4cda97a9ef"}
assert isinstance(results, str)
assert results == "0a1b1bf3-37de-48f7-9863-ed4cda97a9ef"
assert len(events) == 1
assert isinstance(events[0], TriggerEvent)
assert events[0].payload["status"] == "success"
assert events[0].payload["type"] == "builtins.dict"
assert events[0].payload["response"] == json.dumps(status)

@pytest.mark.skipif(not AIRFLOW_V_2_10_PLUS, reason="Lambda parameters works in Airflow >= 2.10.0")
def test_execute_with_lambda_parameter(self):
status = load_json("resources", "status.json")
response = mock_json_response(200, status)

with self.patch_hook_and_request_adapter(response):
sensor = MSGraphSensor(
task_id="check_workspaces_status",
conn_id="powerbi",
url="myorg/admin/workspaces/scanStatus/{scanId}",
path_parameters=lambda context, jinja_env: {"scanId": "0a1b1bf3-37de-48f7-9863-ed4cda97a9ef"},
result_processor=lambda context, result: result["id"],
timeout=350.0,
)

results, events = self.execute_operator(sensor)

assert sensor.path_parameters == {"scanId": "0a1b1bf3-37de-48f7-9863-ed4cda97a9ef"}
assert isinstance(results, str)
assert results == "0a1b1bf3-37de-48f7-9863-ed4cda97a9ef"
assert len(events) == 1
Expand Down
8 changes: 4 additions & 4 deletions providers/tests/microsoft/conftest.py
Original file line number Diff line number Diff line change
Expand Up @@ -132,15 +132,15 @@ def __init__(

def xcom_pull(
self,
task_ids: Iterable[str] | str | None = None,
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,
run_id: str | None = None,
*,
map_indexes: Iterable[int] | int | None = None,
default: Any | None = None,
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}")
Expand Down