fix(compression middleware): Fix resource handling (litestar-org#3927)

euri10 · Jan 5, 2025 · 2fb6f6b · 2fb6f6b
1 parent f5ba3a7
commit 2fb6f6b
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 0 deletions.
diff --git a/litestar/middleware/compression/middleware.py b/litestar/middleware/compression/middleware.py
@@ -131,6 +131,11 @@ async def send_wrapper(message: Message) -> None:
             if initial_message is not None and value_or_default(connection_state.is_cached, False):
                 await send(initial_message)
                 await send(message)
+                facade.close()
+                return
+
+            if initial_message and message["type"] == "http.disconnect":
+                facade.close()
                 return
 
             if initial_message and message["type"] == "http.response.body":
@@ -170,6 +175,7 @@ async def send_wrapper(message: Message) -> None:
                         await send(message)
 
                     else:
+                        facade.close()
                         await send(initial_message)
                         await send(message)
 

diff --git a/tests/unit/test_middleware/test_compression_middleware.py b/tests/unit/test_middleware/test_compression_middleware.py
@@ -196,6 +196,8 @@ async def fake_send(message: Message) -> None:
     # second body message with more_body=True will be empty if zlib buffers output and is not flushed
     await wrapped_send(HTTPResponseBodyEvent(type="http.response.body", body=b"abc", more_body=True))
     assert mock.mock_calls[-1].args[0]["body"]
+    # send a more_body=False so resources close properly
+    await wrapped_send(HTTPResponseBodyEvent(type="http.response.body", body=b"", more_body=False))
 
 
 @pytest.mark.parametrize(