From b8d5927c651c5e049a132021e72f8d01c83d5021 Mon Sep 17 00:00:00 2001 From: RKS Date: Mon, 21 Sep 2026 22:48:59 -0400 Subject: [PATCH] fix: honor Kiota request extensions in Graph transport --- .../middleware/async_graph_transport.py | 11 ++- .../middleware/test_async_graph_transport.py | 77 +++++++++++++++++++ 2 files changed, 86 insertions(+), 2 deletions(-) diff --git a/src/msgraph_core/middleware/async_graph_transport.py b/src/msgraph_core/middleware/async_graph_transport.py index bb81ff69..a6c25e01 100644 --- a/src/msgraph_core/middleware/async_graph_transport.py +++ b/src/msgraph_core/middleware/async_graph_transport.py @@ -6,6 +6,9 @@ from .._enums import FeatureUsageFlag from .request_context import GraphRequestContext +# Older supported Kiota releases do not export this request extension key. +REQUEST_OPTIONS_KEY = 'kiota_request_options' + class AsyncGraphTransport(httpx.AsyncBaseTransport): """A custom transport for requests to the Microsoft Graph API @@ -16,7 +19,9 @@ def __init__(self, transport: httpx.AsyncBaseTransport, pipeline: MiddlewarePipe self.pipeline = pipeline async def handle_async_request(self, request: httpx.Request) -> httpx.Response: - if self.pipeline and hasattr(request, 'options'): + if self.pipeline and ( + REQUEST_OPTIONS_KEY in request.extensions or hasattr(request, 'options') + ): self.set_request_context_and_feature_usage(request) response = await self.pipeline.send(request) return response @@ -26,7 +31,9 @@ async def handle_async_request(self, request: httpx.Request) -> httpx.Response: def set_request_context_and_feature_usage(self, request: httpx.Request) -> httpx.Request: - request_options = request.options # type:ignore + request_options = request.extensions.get(REQUEST_OPTIONS_KEY) + if request_options is None: + request_options = request.options # type:ignore context = GraphRequestContext(request_options, request.headers) middleware = self.pipeline._first_middleware diff --git a/tests/middleware/test_async_graph_transport.py b/tests/middleware/test_async_graph_transport.py index 38e1bca7..025f6495 100644 --- a/tests/middleware/test_async_graph_transport.py +++ b/tests/middleware/test_async_graph_transport.py @@ -1,8 +1,17 @@ +import asyncio + +import httpx import pytest +from kiota_abstractions.authentication import AnonymousAuthenticationProvider +from kiota_abstractions.method import Method +from kiota_abstractions.request_information import RequestInformation +from kiota_http.httpx_request_adapter import HttpxRequestAdapter from kiota_http.kiota_client_factory import KiotaClientFactory from msgraph_core._enums import FeatureUsageFlag +from msgraph_core.graph_client_factory import GraphClientFactory from msgraph_core.middleware import AsyncGraphTransport, GraphRequestContext +from msgraph_core.middleware.async_graph_transport import REQUEST_OPTIONS_KEY def test_set_request_context_and_feature_usage(mock_request, mock_transport): @@ -16,3 +25,71 @@ def test_set_request_context_and_feature_usage(mock_request, mock_transport): assert mock_request.context.feature_usage == hex( FeatureUsageFlag.RETRY_HANDLER_ENABLED | FeatureUsageFlag.REDIRECT_HANDLER_ENABLED ) + + +@pytest.mark.parametrize( + 'content_type', [ + 'application/octet-stream', + 'application/vnd.openxmlformats-officedocument.wordprocessingml.document', + ] +) +def test_binary_download_follows_redirect_with_kiota_request_extensions(content_type): + calls = [] + contexts = [] + + def handle_request(request): + calls.append(str(request.url)) + contexts.append(request.context) + if request.url.host == 'graph.example': + return httpx.Response(302, headers={'Location': 'https://download.example/file'}) + return httpx.Response( + 200, content=b'binary content', headers={'Content-Type': content_type} + ) + + async def download(): + client = GraphClientFactory.create_with_default_middleware( + client=httpx.AsyncClient(transport=httpx.MockTransport(handle_request)) + ) + try: + adapter = HttpxRequestAdapter(AnonymousAuthenticationProvider(), http_client=client) + request_info = RequestInformation() + request_info.http_method = Method.GET + request_info.url = 'https://graph.example/drive/item/content' + return await adapter.send_primitive_async(request_info, 'bytes', {}) + finally: + await client.aclose() + + assert asyncio.run(download()) == b'binary content' + assert calls == ['https://graph.example/drive/item/content', 'https://download.example/file'] + assert all(isinstance(context, GraphRequestContext) for context in contexts) + + +def test_extension_options_take_precedence_over_legacy_attribute(mock_transport): + middleware = KiotaClientFactory.get_default_middleware(None) + pipeline = KiotaClientFactory.create_middleware_pipeline(middleware, mock_transport) + transport = AsyncGraphTransport(mock_transport, pipeline) + request = httpx.Request('GET', 'https://example.org', extensions={REQUEST_OPTIONS_KEY: {}}) + request.options = {'legacy': True} + + transport.set_request_context_and_feature_usage(request) + + assert request.context.middleware_control == {} + + +def test_request_without_options_bypasses_graph_pipeline(): + calls = [] + + def handle_request(request): + calls.append(request) + return httpx.Response(200, content=b'body') + + async def send(): + underlying_transport = httpx.MockTransport(handle_request) + middleware = KiotaClientFactory.get_default_middleware(None) + pipeline = KiotaClientFactory.create_middleware_pipeline(middleware, underlying_transport) + transport = AsyncGraphTransport(underlying_transport, pipeline) + return await transport.handle_async_request(httpx.Request('GET', 'https://example.org')) + + assert asyncio.run(send()).status_code == 200 + assert len(calls) == 1 + assert not hasattr(calls[0], 'context')