diff --git a/CHANGELOG.md b/CHANGELOG.md index 4c56adf..27c105b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,7 +7,14 @@ and adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0). ## Unreleased -Nothing. +## Fixed + +- Fixed `info.response.body` in instrumentation functions being wrongfully empty + if response is not streamed. Affects a feature that was introduced with + release [5.10.0](#5100--2023-02-26) few weeks ago. Closed issue + [#236](https://github.com/trallnag/prometheus-fastapi-instrumentator/issues/236) + and implemented in pull request + [#237](https://github.com/trallnag/prometheus-fastapi-instrumentator/pull/237). ## [5.11.1](https://github.com/trallnag/prometheus-fastapi-instrumentator/compare/v5.11.0...v5.11.1) / 2023-03-11 diff --git a/src/prometheus_fastapi_instrumentator/middleware.py b/src/prometheus_fastapi_instrumentator/middleware.py index 6696c56..df143b9 100644 --- a/src/prometheus_fastapi_instrumentator/middleware.py +++ b/src/prometheus_fastapi_instrumentator/middleware.py @@ -144,9 +144,7 @@ async def send_wrapper(message: Message) -> None: nonlocal status_code, headers headers = message["headers"] status_code = message["status"] - elif message["type"] == "http.response.body" and message.get( - "more_body", False - ): + elif message["type"] == "http.response.body" and message["body"]: nonlocal body body += message["body"] await send(message) diff --git a/tests/test_middleware.py b/tests/test_middleware.py new file mode 100644 index 0000000..4448c1d --- /dev/null +++ b/tests/test_middleware.py @@ -0,0 +1,141 @@ +from fastapi import FastAPI, responses, status +from fastapi.testclient import TestClient + +from prometheus_fastapi_instrumentator import Instrumentator, metrics + + +def test_info_body_empty(): + """ + Tests that `info.response.body` is empty if actual response is also empty. + """ + + app = FastAPI() + client = TestClient(app) + + @app.get("/") + def read_root(): + return responses.Response(status_code=status.HTTP_200_OK) + + instrumentation_executed = False + + def instrumentation(info: metrics.Info) -> None: + nonlocal instrumentation_executed + instrumentation_executed = True + assert len(info.response.body) == 0 + + Instrumentator().instrument(app).add(instrumentation) + + client.get("/") + assert instrumentation_executed + + +def test_info_body_stream_small(): + """ + Tests that `info.response.body` is correct if small response is streamed. + """ + + app = FastAPI() + client = TestClient(app) + + @app.get("/") + def read_root(): + return responses.StreamingResponse((str(num) + "xxx" for num in range(5))) + + instrumentation_executed = False + + def instrumentation(info: metrics.Info) -> None: + nonlocal instrumentation_executed + instrumentation_executed = True + assert len(info.response.body) == 20 + assert info.response.body.decode() == "0xxx1xxx2xxx3xxx4xxx" + + Instrumentator().instrument(app).add(instrumentation) + + response = client.get("/") + assert instrumentation_executed + assert len(response.content) == 20 + assert response.content.decode() == "0xxx1xxx2xxx3xxx4xxx" + + +def test_info_body_stream_large(): + """ + Tests that `info.response.body` is correct if large response is streamed. + """ + + app = FastAPI() + client = TestClient(app) + + @app.get("/") + def read_root(): + return responses.StreamingResponse( + (str(num) + "x" * 10000000 for num in range(5)) + ) + + instrumentation_executed = False + + def instrumentation(info: metrics.Info) -> None: + nonlocal instrumentation_executed + instrumentation_executed = True + assert len(info.response.body) >= 50000000 + + Instrumentator().instrument(app).add(instrumentation) + + response = client.get("/") + assert instrumentation_executed + assert len(response.content) >= 50000000 + + +def test_info_body_bulk_small(): + """ + Tests that `info.response.body` is correct if small response is returned. + """ + + app = FastAPI() + client = TestClient(app) + + @app.get("/", response_class=responses.PlainTextResponse) + def read_root(): + return "123456789" + + instrumentation_executed = False + + def instrumentation(info: metrics.Info) -> None: + print(info.response.body) + nonlocal instrumentation_executed + instrumentation_executed = True + assert len(info.response.body) == 9 + assert info.response.body == b"123456789" + + Instrumentator().instrument(app).add(instrumentation) + + response = client.get("/") + assert instrumentation_executed + assert len(response.content) == 9 + assert response.content == b"123456789" + + +def test_info_body_bulk_large(): + """ + Tests that `info.response.body` is correct if large response is returned. + """ + + app = FastAPI() + client = TestClient(app) + + @app.get("/", response_class=responses.PlainTextResponse) + def read_root(): + return "x" * 50_000_000 + + instrumentation_executed = False + + def instrumentation(info: metrics.Info) -> None: + print(info.response.body) + nonlocal instrumentation_executed + instrumentation_executed = True + assert len(info.response.body) == 50_000_000 + + Instrumentator().instrument(app).add(instrumentation) + + response = client.get("/") + assert instrumentation_executed + assert len(response.content) == 50_000_000