Ensure chunk_size is properly set when chunking responses (#1336)

* Add test case for chunking of an empty response

* Ensure `chunk_size` is set when chunking responses

optimade/server/middleware.py

@@ -439,11 +439,9 @@ async def dispatch(self, request: Request, call_next):
         charset = response.charset
 
         body = b""
-        first_run = True
+        chunk_size = 0
         async for chunk in response.body_iterator:
-            if first_run:
-                first_run = False
-                chunk_size = len(chunk)
+            chunk_size = chunk_size or len(chunk)
             if not isinstance(chunk, bytes):
                 chunk = chunk.encode(charset)
             body += chunk

tests/server/conftest.py

@@ -21,6 +21,16 @@ def index_client():
     return client_factory()(server="index")
 
 
+@pytest.fixture(scope="session", params=["regular"])
+def client_with_empty_extension_endpoint(request):
+    """Return TestClient for the regular OPTIMADE server with an additional
+    empty test endpoint added at `/extensions/test_empty_body`.
+    """
+    from .utils import client_factory
+
+    return client_factory()(server=request.param, add_empty_endpoint=True)
+
+
 @pytest.fixture(scope="session", params=["regular", "index"])
 def both_clients(request):
     """Return TestClient for both the regular and index OPTIMADE server"""

tests/server/middleware/test_warnings.py

@@ -75,3 +75,15 @@ def test_chunk_it_up():
 
     generator = AddWarnings.chunk_it_up(content, chunk_size)
     assert "".join(generator) == content
+
+
+def test_empty_response_chunk_it_up(client_with_empty_extension_endpoint):
+    """Test that the chunking induced by the warnings middleware can handle
+    responses with empty bodies."""
+    from optimade.server.middleware import AddWarnings
+
+    add_warning_middleware = AddWarnings(client_with_empty_extension_endpoint.app)
+
+    response = client_with_empty_extension_endpoint.get("/extensions/test_empty_body")
+    add_warning_middleware._warnings = []
+    assert response.content == b""

tests/server/utils.py

@@ -196,7 +196,19 @@ def inner(
         version: str = None,
         server: str = "regular",
         raise_server_exceptions: bool = True,
+        add_empty_endpoint: bool = False,
     ) -> OptimadeTestClient:
+        """Create a test client for the reference servers parameterised by:
+
+        - `version` (to be prepended to all requests),
+        - `server` type ("regular" or "index"),
+        - whether to raise exceptions from the server to the client
+          (`raise_server_exceptions`),
+        - whether to create an endpoint that returns an empty response
+          at `/extensions/test_empty_body` used for testing streaming
+          responses (`add_empty_endpoint`)
+
+        """
         if server == "regular":
             from optimade.server.main import (
                 app,
@@ -217,6 +229,19 @@ def inner(
         add_major_version_base_url(app)
         add_optional_versioned_base_urls(app)
 
+        if add_empty_endpoint:
+
+            from starlette.routing import Router, Route
+            from fastapi.responses import PlainTextResponse
+
+            async def empty(_):
+                return PlainTextResponse(b"", 200)
+
+            empty_router = Router(
+                routes=[Route("/extensions/test_empty_body", endpoint=empty)]
+            )
+            app.include_router(empty_router)
+
         if version:
             return OptimadeTestClient(
                 app,