diff --git a/packages/http/httpx/kiota_http/middleware/headers_inspection_handler.py b/packages/http/httpx/kiota_http/middleware/headers_inspection_handler.py index f03facf1..c86e7ad1 100644 --- a/packages/http/httpx/kiota_http/middleware/headers_inspection_handler.py +++ b/packages/http/httpx/kiota_http/middleware/headers_inspection_handler.py @@ -6,6 +6,7 @@ from typing import Optional +from kiota_abstractions.headers_collection import HeadersCollection from kiota_abstractions.request_option import RequestOption import httpx @@ -46,19 +47,36 @@ async def send( Returns: Response: The response object. """ + if request is None: + raise TypeError("request cannot be null") + current_options = self._get_current_options(request) span = self._create_observability_span(request, "HeadersInspectionHandler_send") - span.set_attribute(HEADERS_INSPECTION_KEY, True) - span.end() - - if current_options and current_options.inspect_request_headers: - for header in request.headers: - current_options.request_headers.add(header, request.headers[header]) - response = await super().send(request, transport) - if current_options and current_options.inspect_response_headers: - for header in response.headers: - current_options.response_headers.add(header, response.headers[header]) - return response + try: + span.set_attribute(HEADERS_INSPECTION_KEY, True) + + if current_options and current_options.inspect_request_headers: + self._capture_headers(request.headers, current_options.request_headers) + response = await super().send(request, transport) + if current_options and current_options.inspect_response_headers: + self._capture_headers(response.headers, current_options.response_headers) + return response + finally: + span.end() + + def _capture_headers( + self, source_headers: httpx.Headers, target_collection: HeadersCollection + ) -> None: + """Captures headers from httpx.Headers into a HeadersCollection. + + Headers belonging to SINGLE_VALUE_HEADERS preserve their combined string value, + while other multi-valued headers preserve distinct list items. + """ + for header in source_headers: + if header.lower() in HeadersCollection.SINGLE_VALUE_HEADERS: + target_collection.add(header, source_headers[header]) + else: + target_collection.add(header, source_headers.get_list(header)) def _get_current_options(self, request: httpx.Request) -> HeadersInspectionHandlerOption: """Returns the options to use for the request.Overrides default options if @@ -74,10 +92,10 @@ def _get_current_options(self, request: httpx.Request) -> HeadersInspectionHandl request_options = request.extensions.get(REQUEST_OPTIONS_KEY) if request_options: current_options = request_options.get(HeadersInspectionHandlerOption.get_key(), None) - if current_options: - return current_options + if not current_options: + current_options = self.options # Clear headers per request - self.options.request_headers.clear() - self.options.response_headers.clear() - return self.options + current_options.request_headers.clear() + current_options.response_headers.clear() + return current_options diff --git a/packages/http/httpx/tests/middleware_tests/test_headers_inspection_handler.py b/packages/http/httpx/tests/middleware_tests/test_headers_inspection_handler.py index d0704140..2859de96 100644 --- a/packages/http/httpx/tests/middleware_tests/test_headers_inspection_handler.py +++ b/packages/http/httpx/tests/middleware_tests/test_headers_inspection_handler.py @@ -1,12 +1,12 @@ import pytest -import httpx - -from unittest.mock import AsyncMock - from kiota_abstractions.headers_collection import HeadersCollection -from kiota_http.middleware.middleware import BaseMiddleware -from kiota_http.middleware.options.headers_inspection_handler_option import HeadersInspectionHandlerOption + +import httpx from kiota_http.middleware.headers_inspection_handler import HeadersInspectionHandler +from kiota_http.middleware.middleware import REQUEST_OPTIONS_KEY, BaseMiddleware +from kiota_http.middleware.options.headers_inspection_handler_option import ( + HeadersInspectionHandlerOption, +) def test_no_config(): @@ -95,3 +95,133 @@ def request_handler(request: httpx.Request): assert not handler.options.request_headers.try_get('test_request') == {'test_request_header'} assert handler.options.request_headers.try_get('test_request_2') == {'test_request_header_2'} assert handler.options.response_headers.try_get('test_response') == {'test_response_header'} + + +@pytest.mark.asyncio +async def test_send_request_none_raises_type_error(): + handler = HeadersInspectionHandler() + mock_transport = httpx.MockTransport(lambda req: httpx.Response(200)) + with pytest.raises(TypeError, match="request cannot be null"): + await handler.send(None, mock_transport) + + +@pytest.mark.asyncio +async def test_headers_inspection_handler_request_options_isolation_and_clearing(): + + def request_handler(request: httpx.Request): + return httpx.Response(200, headers={"test_response": "resp_val_1"}) + + handler = HeadersInspectionHandler() + custom_options = HeadersInspectionHandlerOption() + + request1 = httpx.Request( + "GET", + "https://localhost", + headers={"test_request": "req_val_1"}, + extensions={ + REQUEST_OPTIONS_KEY: { + HeadersInspectionHandlerOption.get_key(): custom_options + } + }, + ) + mock_transport = httpx.MockTransport(request_handler) + resp1 = await handler.send(request1, mock_transport) + assert resp1.status_code == 200 + assert custom_options.request_headers.try_get("test_request") == {"req_val_1"} + assert custom_options.response_headers.try_get("test_response") == {"resp_val_1"} + + # Second request reusing custom_options should clear previous request/response headers + def request_handler_2(request: httpx.Request): + return httpx.Response(200, headers={"test_response_2": "resp_val_2"}) + + request2 = httpx.Request( + "GET", + "https://localhost", + headers={"test_request_2": "req_val_2"}, + extensions={ + REQUEST_OPTIONS_KEY: { + HeadersInspectionHandlerOption.get_key(): custom_options + } + }, + ) + mock_transport_2 = httpx.MockTransport(request_handler_2) + resp2 = await handler.send(request2, mock_transport_2) + assert resp2.status_code == 200 + assert custom_options.request_headers.try_get("test_request") is False + assert custom_options.request_headers.try_get("test_request_2") == {"req_val_2"} + assert custom_options.response_headers.try_get("test_response") is False + assert custom_options.response_headers.try_get("test_response_2") == {"resp_val_2"} + + +@pytest.mark.asyncio +async def test_headers_inspection_handler_multi_valued_headers(): + + def request_handler(request: httpx.Request): + return httpx.Response( + 200, + headers=[ + ("Set-Cookie", "session=abc; Path=/"), + ("Set-Cookie", "theme=dark; Path=/"), + ("X-Custom-List", "item1"), + ("X-Custom-List", "item2"), + ("Content-Encoding", "gzip"), + ("Content-Encoding", "br"), + ], + ) + + handler = HeadersInspectionHandler() + request = httpx.Request( + "GET", + "https://localhost", + headers=[ + ("Accept", "application/json"), + ("Accept", "text/plain"), + ("Content-Encoding", "gzip"), + ("Content-Encoding", "br"), + ], + ) + mock_transport = httpx.MockTransport(request_handler) + resp = await handler.send(request, mock_transport) + assert resp.status_code == 200 + assert handler.options.request_headers.try_get("accept") == {"application/json", "text/plain"} + assert handler.options.request_headers.try_get("content-encoding") == {"gzip, br"} + assert handler.options.response_headers.try_get("set-cookie") == { + "session=abc; Path=/", + "theme=dark; Path=/", + } + assert handler.options.response_headers.try_get("x-custom-list") == {"item1", "item2"} + assert handler.options.response_headers.try_get("content-encoding") == {"gzip, br"} + + +@pytest.mark.asyncio +async def test_span_ends_on_transport_failure(span_exporter): + + def request_handler(request): + assert len(span_exporter.get_finished_spans()) == 0 + raise httpx.ReadError("request interrupted") + + handler = HeadersInspectionHandler() + request = httpx.Request("GET", "https://localhost") + async with httpx.MockTransport(request_handler) as transport: + with pytest.raises(httpx.ReadError): + await handler.send(request, transport) + + assert [span.name + for span in span_exporter.get_finished_spans()] == ["HeadersInspectionHandler_send"] + + +@pytest.mark.asyncio +async def test_span_remains_open_during_successful_transport(span_exporter): + + def request_handler(request): + assert len(span_exporter.get_finished_spans()) == 0 + return httpx.Response(200) + + handler = HeadersInspectionHandler() + request = httpx.Request("GET", "https://localhost") + async with httpx.MockTransport(request_handler) as transport: + resp = await handler.send(request, transport) + + assert resp.status_code == 200 + assert [span.name + for span in span_exporter.get_finished_spans()] == ["HeadersInspectionHandler_send"]