Cancel task group after returning middleware response

Add test for encode#1022
uSpike · Jun 13, 2021 · 63cfcb9 · 63cfcb9
1 parent a1ceb35
commit 63cfcb9
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 0 deletions.
diff --git a/starlette/middleware/base.py b/starlette/middleware/base.py
@@ -54,6 +54,7 @@ async def body_stream() -> typing.AsyncGenerator[bytes, None]:
             request = Request(scope, receive=receive)
             response = await self.dispatch_func(request, call_next)
             await response(scope, receive, send)
+            task_group.cancel_scope.cancel()
 
     async def dispatch(
         self, request: Request, call_next: RequestResponseEndpoint

diff --git a/tests/middleware/test_base.py b/tests/middleware/test_base.py
@@ -143,3 +143,18 @@ def homepage(request):
 def test_middleware_repr():
     middleware = Middleware(CustomMiddleware)
     assert repr(middleware) == "Middleware(CustomMiddleware)"
+
+
+def test_fully_evaluated_response():
+    # Test for https://github.com/encode/starlette/issues/1022
+    class CustomMiddleware(BaseHTTPMiddleware):
+        async def dispatch(self, request, call_next):
+            await call_next(request)
+            return PlainTextResponse("Custom")
+
+    app = Starlette()
+    app.add_middleware(CustomMiddleware)
+
+    client = TestClient(app)
+    response = client.get("/does_not_exist")
+    assert response.text == "Custom"