From 2fb6f6b118f1247421db027fafb14507b3f8f84d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Janek=20Nouvertn=C3=A9?= Date: Sun, 5 Jan 2025 17:13:45 +0100 Subject: [PATCH] fix(compression middleware): Fix resource handling (#3927) --- litestar/middleware/compression/middleware.py | 6 ++++++ tests/unit/test_middleware/test_compression_middleware.py | 2 ++ 2 files changed, 8 insertions(+) diff --git a/litestar/middleware/compression/middleware.py b/litestar/middleware/compression/middleware.py index 7ea7853b08..5458bcc8d7 100644 --- 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 index 7c5f04ec07..1feb001bb5 100644 --- 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(