From af0e946ff8824a36ba919ac77b0ee0e3cfb2f79f Mon Sep 17 00:00:00 2001 From: Jarkko Jaakola Date: Wed, 29 Jan 2025 11:52:36 +0200 Subject: [PATCH] fix: re-raise exception in telemetry middleware --- src/schema_registry/telemetry/middleware.py | 7 ++- .../telemetry/test_middleware.py | 53 ++++++++++--------- 2 files changed, 31 insertions(+), 29 deletions(-) diff --git a/src/schema_registry/telemetry/middleware.py b/src/schema_registry/telemetry/middleware.py index bc6788d89..5a6041d3b 100644 --- a/src/schema_registry/telemetry/middleware.py +++ b/src/schema_registry/telemetry/middleware.py @@ -35,15 +35,14 @@ async def telemetry_middleware( http_request_metrics.record_request_exception(ATTRIBUTES=ATTRIBUTES, exc=exc) SPAN.set_status(Status(StatusCode.ERROR)) SPAN.record_exception(exc) + http_request_metrics.finish_request(ATTRIBUTES=ATTRIBUTES, request=request, response=None) + raise exc else: tracer.update_span_with_response(response=response, span=SPAN) - return response - finally: http_request_metrics.finish_request( ATTRIBUTES=ATTRIBUTES, request=request, response=response if "response" in locals() else None ) - - return None + return response def setup_telemetry_middleware(app: FastAPI) -> None: diff --git a/tests/unit/schema_registry/telemetry/test_middleware.py b/tests/unit/schema_registry/telemetry/test_middleware.py index 60647aa10..6b50ed1f8 100644 --- a/tests/unit/schema_registry/telemetry/test_middleware.py +++ b/tests/unit/schema_registry/telemetry/test_middleware.py @@ -104,29 +104,32 @@ async def test_telemetry_middleware_call_next_exception(http_request_metrics: Ma SpanStatus = MagicMock(spec=Status, status_code=StatusCode.ERROR) + response = None with patch("schema_registry.telemetry.middleware.Status", return_value=SpanStatus): - response = await telemetry_middleware( - request=request_mock, call_next=call_next, tracer=tracer, http_request_metrics=http_request_metrics - ) - span = tracer.get_tracer.return_value.start_as_current_span.return_value.__enter__.return_value - - tracer.get_tracer.assert_called_once() - tracer.get_tracer.return_value.start_as_current_span.assert_called_once_with(name="GET: /test", kind=SpanKind.SERVER) - tracer.update_span_with_request.assert_called_once_with(request=request_mock, span=span) - - # Check that the request handler is called - call_next.assert_awaited_once_with(request_mock) - span.set_status.assert_called_once_with(SpanStatus) - - http_request_metrics.assert_has_calls( - [ - call.get_resource_from_request(request=request_mock), - call.start_request(request=request_mock), - call.record_request_exception(ATTRIBUTES=http_request_metrics.start_request.return_value, exc=exception), - call.finish_request( - ATTRIBUTES=http_request_metrics.start_request.return_value, request=request_mock, response=None - ), - ] - ) - - assert not response + with pytest.raises(Exception) as excinfo: + response = await telemetry_middleware( + request=request_mock, call_next=call_next, tracer=tracer, http_request_metrics=http_request_metrics + ) + + assert "Test exception" in str(excinfo.value) + assert response is None + span = tracer.get_tracer.return_value.start_as_current_span.return_value.__enter__.return_value + + tracer.get_tracer.assert_called_once() + tracer.get_tracer.return_value.start_as_current_span.assert_called_once_with(name="GET: /test", kind=SpanKind.SERVER) + tracer.update_span_with_request.assert_called_once_with(request=request_mock, span=span) + + # Check that the request handler is called + call_next.assert_awaited_once_with(request_mock) + span.set_status.assert_called_once_with(SpanStatus) + + http_request_metrics.assert_has_calls( + [ + call.get_resource_from_request(request=request_mock), + call.start_request(request=request_mock), + call.record_request_exception(ATTRIBUTES=http_request_metrics.start_request.return_value, exc=exception), + call.finish_request( + ATTRIBUTES=http_request_metrics.start_request.return_value, request=request_mock, response=None + ), + ] + )