fix: Fix info.response.body empty with normal responses (#237)

- chore: Fix
- chore: Add tests
- chore: Update changelog

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
 

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)

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