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 4590a361026..40420bebee4 100644 --- a/tests/test_litellm/llms/custom_httpx/test_http_handler.py +++ b/tests/test_litellm/llms/custom_httpx/test_http_handler.py @@ -1034,27 +1034,56 @@ def _mock_transport() -> httpx.MockTransport: return httpx.MockTransport(respond) +RELEASED_TOO_EARLY = "the handler was released while its response could still read" +NEVER_RELEASED = "the handler outlived the response that was holding it" + +# Every method that can hand back a body the caller has not read yet, which is +# every one that passes stream= down to send(). Parametrized so a method added +# later is covered here rather than being the one that forgets to anchor. +ASYNC_STREAMING_SENDS = ["post", "delete"] +SYNC_STREAMING_SENDS = ["post", "patch", "put", "delete"] + + @pytest.mark.asyncio -async def test_a_streaming_response_holds_its_handler_until_it_is_released(): +@pytest.mark.parametrize("method", ASYNC_STREAMING_SENDS) +async def test_a_streaming_response_holds_its_handler_until_it_is_released(method): """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. + withholds the close, and releasing the anchor 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) + response = await getattr(handler, method)("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 ref() is not None, RELEASED_TOO_EARLY assert await response.aread() == b"ab" del response gc.collect() - assert ref() is None, "the handler outlived the response that was holding it" + assert ref() is None, NEVER_RELEASED + + +@pytest.mark.parametrize("method", SYNC_STREAMING_SENDS) +def test_a_sync_streaming_response_holds_its_handler_until_it_is_released(method): + """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 = getattr(handler, method)("https://example.invalid/stream", stream=True) + + del handler + gc.collect() + assert ref() is not None, RELEASED_TOO_EARLY + + assert response.read() == b"ab" + del response + gc.collect() + assert ref() is None, NEVER_RELEASED @pytest.mark.asyncio @@ -1076,23 +1105,6 @@ async def test_a_fully_read_response_does_not_hold_its_handler(): 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)