fix(http): handle null request, keep span open over call, and isolate options in HeadersInspectionHandler

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

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"]