From 8cf17717d0e5cd2172461c5c3c4f26c518abd8dc Mon Sep 17 00:00:00 2001 From: Luke Maisel Date: Wed, 20 Aug 2025 16:51:53 -0400 Subject: [PATCH] fix: ensures BaseMiddleware always creates a new span instead of returning the current span Updated _create_observability_span to create a new span with the middleware's parent_span context or current context if needed. Updated BaseMiddleware tests to be more accurate and ensure BaseMiddleware does not end outer spans. --- .../httpx/kiota_http/middleware/middleware.py | 8 +++++--- packages/http/httpx/tests/conftest.py | 5 +++++ .../middleware_tests/test_base_middleware.py | 18 +++++++++++++++--- 3 files changed, 25 insertions(+), 6 deletions(-) diff --git a/packages/http/httpx/kiota_http/middleware/middleware.py b/packages/http/httpx/kiota_http/middleware/middleware.py index 85075158..6e70e385 100644 --- a/packages/http/httpx/kiota_http/middleware/middleware.py +++ b/packages/http/httpx/kiota_http/middleware/middleware.py @@ -64,8 +64,9 @@ async def send(self, request, transport): return await self.next.send(request, transport) def _create_observability_span(self, request, span_name: str) -> trace.Span: - """Gets the parent_span from the request options and creates a new span. - If no parent_span is found, we try to get the current span.""" + """Gets the parent_span from the request options and creates a new span. + If no parent_span is found in the request, uses the parent_span in the + object. If parent_span is None, current context will be used.""" _span = None if options := getattr(request, "options", None): if parent_span := options.get("parent_span", None): @@ -73,5 +74,6 @@ def _create_observability_span(self, request, span_name: str) -> trace.Span: _context = trace.set_span_in_context(parent_span) _span = tracer.start_span(span_name, _context) if _span is None: - _span = trace.get_current_span() + _context = trace.set_span_in_context(self.parent_span) + _span = tracer.start_span(span_name, _context) return _span diff --git a/packages/http/httpx/tests/conftest.py b/packages/http/httpx/tests/conftest.py index c008898f..55c9c124 100644 --- a/packages/http/httpx/tests/conftest.py +++ b/packages/http/httpx/tests/conftest.py @@ -211,6 +211,11 @@ def mock_no_content_response(sample_headers): tracer = trace.get_tracer(__name__) +@pytest.fixture +def otel_tracer(): + return tracer + + @pytest.fixture def mock_otel_span(): return tracer.start_span("mock") diff --git a/packages/http/httpx/tests/middleware_tests/test_base_middleware.py b/packages/http/httpx/tests/middleware_tests/test_base_middleware.py index 254ecabc..2f9ed794 100644 --- a/packages/http/httpx/tests/middleware_tests/test_base_middleware.py +++ b/packages/http/httpx/tests/middleware_tests/test_base_middleware.py @@ -10,9 +10,21 @@ def test_next_is_none(): assert middleware.next is None -def test_span_created(request_info): - """Ensures the current span is returned and the parent_span is not set.""" +def test_span_returned(request_info): + """Ensures a span is returned and the parent_span is not set.""" middleware = BaseMiddleware() - span = middleware._create_observability_span(request_info, "test_span_created") + span = middleware._create_observability_span(request_info, "test_span_returned") + span.end() assert isinstance(span, trace.Span) assert middleware.parent_span is None + + +def test_span_created(request_info, otel_tracer): + """Ensures the span returned does not end an outer context managed span""" + with otel_tracer.start_as_current_span("outside-span") as outside_span: + before_state = outside_span.is_recording() + middleware = BaseMiddleware() + span = middleware._create_observability_span(request_info, "test_span_created") + span.end() + after_state = outside_span.is_recording() + assert(before_state == after_state) \ No newline at end of file