diff --git a/packages/python-sdk/tests/test_file_stream_reader.py b/packages/python-sdk/tests/test_file_stream_reader.py index f15cedda2b..80a7a4933c 100644 --- a/packages/python-sdk/tests/test_file_stream_reader.py +++ b/packages/python-sdk/tests/test_file_stream_reader.py @@ -1,12 +1,15 @@ """Unit tests for the streamed-read helpers. These exercise the readers' own lifecycle (consume / context manager / -explicit close / idle timeout) against a local chunked HTTP server. They -assert on the reader's contract — the underlying response is closed — rather -than on connection-pool internals, which are private to the transport and -absent on the pyqwest transports the SDK ships. +explicit close / read error / abandonment / idle timeout) against a local +chunked HTTP server. They assert on the reader's contract — the underlying +response is closed, or deliberately left open — rather than on +connection-pool internals, which are private to the transport and absent on +the pyqwest transports the SDK ships. """ +import asyncio +import gc import socket import threading import time @@ -24,15 +27,32 @@ EXPECTED = b"".join(CHUNKS) +def _read_request_head(conn) -> None: + """Read up to the end of the request head. + + Accumulates across reads, since the head can arrive split across TCP + segments, and stops on a closed peer rather than spinning on empty reads. + """ + buffered = b"" + while b"\r\n\r\n" not in buffered: + received = conn.recv(65536) + if not received: + return + buffered += received + + def _start_chunked_server( stall_before: Optional[int] = None, stall_seconds: float = 0.0, + truncate_before: Optional[int] = None, ) -> int: """Start a one-shot HTTP server that replies with a chunked body. When ``stall_before`` is not None, the server sleeps ``stall_seconds`` before sending that chunk index, so a reader with a shorter idle timeout - times out. Returns the server's port. + times out. When ``truncate_before`` is not None, the server drops the + connection at that chunk index without the terminating zero-length chunk, + so a mid-body read raises a protocol error. Returns the server's port. """ sock = socket.socket() sock.bind(("127.0.0.1", 0)) @@ -42,49 +62,20 @@ def _start_chunked_server( def serve(): try: conn, _ = sock.accept() - while b"\r\n\r\n" not in conn.recv(65536): - pass + _read_request_head(conn) conn.sendall( b"HTTP/1.1 200 OK\r\n" b"Content-Type: application/octet-stream\r\n" b"Transfer-Encoding: chunked\r\n\r\n" ) for idx, chunk in enumerate(CHUNKS): + if idx == truncate_before: + break if idx == stall_before: time.sleep(stall_seconds) conn.sendall(f"{len(chunk):x}\r\n".encode() + chunk + b"\r\n") - conn.sendall(b"0\r\n\r\n") - conn.close() - except OSError: - pass - finally: - sock.close() - - threading.Thread(target=serve, daemon=True).start() - return port - - -def _start_truncating_server() -> int: - """One-shot server that sends the head and one chunk, then drops the - connection without the terminating zero-length chunk, so a mid-body read - raises a protocol error.""" - sock = socket.socket() - sock.bind(("127.0.0.1", 0)) - sock.listen(1) - port = sock.getsockname()[1] - - def serve(): - try: - conn, _ = sock.accept() - while b"\r\n\r\n" not in conn.recv(65536): - pass - conn.sendall( - b"HTTP/1.1 200 OK\r\n" - b"Content-Type: application/octet-stream\r\n" - b"Transfer-Encoding: chunked\r\n\r\n" - ) - chunk = CHUNKS[0] - conn.sendall(f"{len(chunk):x}\r\n".encode() + chunk + b"\r\n") + else: + conn.sendall(b"0\r\n\r\n") conn.close() except OSError: pass @@ -102,7 +93,7 @@ def _open_stream(client, port): return client.send(request, stream=True) -def test_sync_full_consume_releases_connection(): +def test_sync_full_consume_releases_response(): with httpx.Client() as client: port = _start_chunked_server() response = _open_stream(client, port) @@ -135,18 +126,37 @@ def test_sync_close_is_idempotent(): def test_sync_read_error_releases_response(): with httpx.Client() as client: - port = _start_truncating_server() + port = _start_chunked_server(truncate_before=1) response = _open_stream(client, port) reader = FileStreamReader(response) it = iter(reader) assert next(it) == CHUNKS[0] # A mid-body error propagates and the reader releases the response. - with pytest.raises(httpx.RemoteProtocolError): + # Asserted on the `httpx.HTTPError` base rather than the concrete + # class, because which error a truncated body surfaces as is the + # transport's business, not part of the reader's contract. + with pytest.raises(httpx.HTTPError): next(it) assert response.is_closed -async def test_async_full_consume_releases_connection(): +def test_sync_abandoned_reader_leaves_the_response_open(): + with httpx.Client() as client: + port = _start_chunked_server() + response = _open_stream(client, port) + reader = FileStreamReader(response) + assert next(iter(reader)) == CHUNKS[0] + + # The reader documents that it has no garbage-collection safety net, + # so dropping it half-consumed must leave the response — and the + # pooled connection behind it — open. Callers have to consume it + # fully, use the context manager, or call `close()`. + del reader + gc.collect() + assert not response.is_closed + + +async def test_async_full_consume_releases_response(): async with httpx.AsyncClient() as client: port = _start_chunked_server() request = client.build_request("GET", f"http://127.0.0.1:{port}/files") @@ -179,6 +189,25 @@ async def test_async_aclose_is_idempotent(): assert response.is_closed +async def test_async_abandoned_reader_leaves_the_response_open(): + async with httpx.AsyncClient() as client: + port = _start_chunked_server() + request = client.build_request("GET", f"http://127.0.0.1:{port}/files") + response = await client.send(request, stream=True) + reader = AsyncFileStreamReader(response) + assert await reader.__anext__() == CHUNKS[0] + + # Same contract as the sync reader, and for a stronger reason: + # releasing an async response requires awaiting `aclose()`, which a + # finalizer cannot do. Sleeping gives the event loop a chance to run + # async-generator finalization, so this asserts the response really + # stays open rather than just outracing the loop. + del reader + gc.collect() + await asyncio.sleep(0.05) + assert not response.is_closed + + async def test_async_reader_explicit_idle_timeout_bounds_each_read(): # The per-call idle bound is enforced with wait_for around each read, so # it works on the regular transport (no transport-level read timeout). diff --git a/packages/python-sdk/tests/test_volume_client.py b/packages/python-sdk/tests/test_volume_client.py index 6429daca2e..28f9446597 100644 --- a/packages/python-sdk/tests/test_volume_client.py +++ b/packages/python-sdk/tests/test_volume_client.py @@ -169,6 +169,20 @@ async def get_transports(): CHUNK = b"x" * 1024 +def _read_request_head(conn) -> None: + """Read up to the end of the request head. + + Accumulates across reads, since the head can arrive split across TCP + segments, and stops on a closed peer rather than spinning on empty reads. + """ + buffered = b"" + while b"\r\n\r\n" not in buffered: + received = conn.recv(65536) + if not received: + return + buffered += received + + def _start_volume_file_server( chunk_delays: List[float], ttfb_delay: float = 0.0 ) -> str: @@ -183,8 +197,7 @@ def _start_volume_file_server( def serve(): try: conn, _ = sock.accept() - while b"\r\n\r\n" not in conn.recv(65536): - pass + _read_request_head(conn) time.sleep(ttfb_delay) conn.sendall( b"HTTP/1.1 200 OK\r\n" @@ -244,6 +257,38 @@ def test_sync_stream_stall_raises_read_timeout(short_read_timeout): assert received == [CHUNK] +# A body far past the ~500 KiB the socket and reqwest buffers hold between +# them, so a stall a quarter of the way in provably lands mid-transfer with the +# server still pushing, rather than after the whole body was read ahead. +SLOW_CONSUMER_CHUNKS = 2048 +SLOW_CONSUMER_STALL_AT = SLOW_CONSUMER_CHUNKS // 4 + + +def test_sync_stream_survives_a_consumer_slower_than_read_timeout(short_read_timeout): + # The idle read timeout is wire-only: the server sends every chunk + # promptly and the consumer then pauses for far longer than the bound, + # which must not abort the stream (parity with the JS SDK's + # "wrapStreamWithConnectionCleanup does not abort a slow consumer"). + # This guards the bound staying an idle one — putting a wall-clock deadline + # back on the streaming path would fail here while every stall test above + # kept passing. + api_url = _start_volume_file_server([0.0] * SLOW_CONSUMER_CHUNKS) + volume = Volume(volume_id="v1", name="test", token="vol-token") + + stream = volume.read_file("file.bin", format="stream", api_url=api_url) + received = 0 + stalled = False + for chunk in stream: + received += len(chunk) + # A threshold rather than an exact count: the transport is free to + # coalesce or split the chunks it hands back. + if not stalled and received >= len(CHUNK) * SLOW_CONSUMER_STALL_AT: + stalled = True + time.sleep(short_read_timeout * 3) + assert stalled + assert received == len(CHUNK) * SLOW_CONSUMER_CHUNKS + + def test_async_stream_survives_transfers_longer_than_read_timeout(short_read_timeout): api_url = _start_volume_file_server([0.15] * 4) volume = AsyncVolume(volume_id="v1", name="test", token="vol-token") @@ -255,6 +300,54 @@ async def run(): assert asyncio.run(run()) == CHUNK * 4 +def test_async_stream_survives_a_consumer_slower_than_read_timeout(short_read_timeout): + # Async counterpart of the sync slow-consumer case above, against the + # async streaming transport. + api_url = _start_volume_file_server([0.0] * SLOW_CONSUMER_CHUNKS) + volume = AsyncVolume(volume_id="v1", name="test", token="vol-token") + + async def run(): + stream = await volume.read_file("file.bin", format="stream", api_url=api_url) + received = 0 + stalled = False + async for chunk in stream: + received += len(chunk) + if not stalled and received >= len(CHUNK) * SLOW_CONSUMER_STALL_AT: + stalled = True + await asyncio.sleep(short_read_timeout * 3) + return stalled, received + + assert asyncio.run(run()) == (True, len(CHUNK) * SLOW_CONSUMER_CHUNKS) + + +def test_sync_stream_explicit_request_timeout_is_a_total_transfer_deadline( + short_read_timeout, +): + # The counterpart that makes the two tests above meaningful: same fixture, + # same body, same stall — the only difference is the explicit + # `request_timeout`, which is documented as a whole-transfer deadline and + # so *does* cut off a slow consumer. Without this A/B, a regression that + # turned the default idle bound into a wall-clock one would look + # indistinguishable from correct behavior. The stall is longer than the + # total deadline, and the tests above establish that a slow consumer never + # trips the idle bound, so the timeout here can only be the deadline. + api_url = _start_volume_file_server([0.0] * SLOW_CONSUMER_CHUNKS) + volume = Volume(volume_id="v1", name="test", token="vol-token") + + stream = volume.read_file( + "file.bin", format="stream", request_timeout=1.0, api_url=api_url + ) + received = 0 + stalled = False + with pytest.raises(httpx.TimeoutException): + for chunk in stream: + received += len(chunk) + if not stalled and received >= len(CHUNK) * SLOW_CONSUMER_STALL_AT: + stalled = True + time.sleep(2.0) + assert stalled + + def test_async_stream_stall_raises_read_timeout(short_read_timeout): api_url = _start_volume_file_server([0.0, 5.0]) volume = AsyncVolume(volume_id="v1", name="test", token="vol-token")