diff --git a/tests/test_litellm/llms/custom_httpx/test_http_handler.py b/tests/test_litellm/llms/custom_httpx/test_http_handler.py index 9e64bfafa54..4590a361026 100644 --- a/tests/test_litellm/llms/custom_httpx/test_http_handler.py +++ b/tests/test_litellm/llms/custom_httpx/test_http_handler.py @@ -1025,6 +1025,74 @@ def test_handed_out_sync_client_pool_survives_handler_collection(keepalive_serve consumer_client.close() +def _mock_transport() -> httpx.MockTransport: + """Answers anything with a short body, left unread when the caller asked to stream.""" + + def respond(request: httpx.Request) -> httpx.Response: + return httpx.Response(200, request=request, content=b"ab") + + return httpx.MockTransport(respond) + + +@pytest.mark.asyncio +async def test_a_streaming_response_holds_its_handler_until_it_is_released(): + """The finalizer must not run while a body this handler issued can still arrive. + + ``_handler_may_close_client`` cannot see that body: it holds the connection it + reads from and never the client. Anchoring the handler to the response is what + withholds the close, and releasing it is what still delivers one. + """ + handler = AsyncHTTPHandler() + handler.client._transport = _mock_transport() + ref = weakref.ref(handler) + response = await handler.post("https://example.invalid/stream", stream=True) + + del handler + gc.collect() + assert ref() is not None, "the handler was released while its response could still read" + + assert await response.aread() == b"ab" + del response + gc.collect() + assert ref() is None, "the handler outlived the response that was holding it" + + +@pytest.mark.asyncio +async def test_a_fully_read_response_does_not_hold_its_handler(): + """A non-streaming response is complete when ``post`` returns, so it anchors nothing. + + Otherwise every client close would wait on whatever the caller does next with + a response it has already read. + """ + handler = AsyncHTTPHandler() + handler.client._transport = _mock_transport() + ref = weakref.ref(handler) + response = await handler.post("https://example.invalid/whole") + assert response.content == b"ab" + + del handler + gc.collect() + + assert ref() is None, "a fully-read response pinned its handler" + + +def test_a_sync_streaming_response_holds_its_handler_until_it_is_released(): + """The sync finalizer closes inline, so the same anchor has to hold it off.""" + handler = HTTPHandler() + handler.client._transport = _mock_transport() + ref = weakref.ref(handler) + response = handler.post("https://example.invalid/stream", stream=True) + + del handler + gc.collect() + assert ref() is not None, "the handler was released while its response could still read" + + assert response.read() == b"ab" + del response + gc.collect() + assert ref() is None, "the handler outlived the response that was holding it" + + def test_sync_close_leaves_caller_supplied_client_open(): supplied = httpx.Client() handler = HTTPHandler(client=supplied)