Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
115 changes: 72 additions & 43 deletions packages/python-sdk/tests/test_file_stream_reader.py
Original file line number Diff line number Diff line change
@@ -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
Expand All @@ -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))
Expand All @@ -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
Expand All @@ -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)
Expand Down Expand Up @@ -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.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Low priority, and the stated reasoning is sound — this file drives the readers against a bare httpx.Client(), so which concrete error a truncated body produces genuinely is not the reader's contract.

Worth knowing what the widening costs, though: httpx.ReadTimeout and httpx.TimeoutException are both subclasses of httpx.HTTPError, so after this change the test no longer distinguishes "the peer truncated the body and we noticed promptly" from "we hung until some idle bound expired". Both close the response, so the is_closed assertion does not separate them either.

The release behaviour itself is still well pinned — removing the except BaseException: self.close() from __next__ fails this test and only this test (the full-consume test survives it, since httpx closes on a clean end-of-stream regardless). If you want the promptness back without re-coupling to the transport's error taxonomy, asserting the elapsed time is well under the idle bound would do it in one line.

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")
Expand Down Expand Up @@ -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).
Expand Down
97 changes: 95 additions & 2 deletions packages/python-sdk/tests/test_volume_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -169,6 +169,20 @@ async def get_transports():
CHUNK = b"x" * 1024


def _read_request_head(conn) -> None:

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nit: this helper is byte-identical to the one added at test_file_stream_reader.py:30. The description's stated reason for fixing the loop at all is that #1690 copied it into a second helper — and the fix lands as a third and fourth copy rather than one shared one.

Sharing is already set up: pytest.ini sets pythonpath = tests with the comment "Makes shared test helpers (e.g. envd_frame_server) importable under importlib mode", and this file already imports transport_caches. A tests/http_server.py (or an addition to an existing helper module) would leave one place for the next person to fix.

"""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:
Expand All @@ -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"
Expand Down Expand Up @@ -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
Comment thread
mishushakov marked this conversation as resolved.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The finding worth acting on. These two tests are non-vacuous on an idle machine, but they degenerate under the load CI actually runs, and the mechanism that makes them work is not the one this comment describes.

I reproduced the fixture verbatim — the time.sleep(delay) call matters, because time.sleep(0) yields the GIL — and recorded server progress at the instant the consumer stalls.

On an idle VM the property holds: the server was at 512–572 of 2048 chunks with body_complete=False in 8/8 runs (4 sync, 4 async). But note 512 of 2048 — that is exactly the consumer's own read position, i.e. essentially zero read-ahead. What holds the invariant is that the server thread, yielding the GIL on every time.sleep(0), is no faster than the consumer. It is a scheduling coincidence, not buffer capacity.

The capacity claim is off by more than an order of magnitude, in the direction that removes the headroom. With the GIL yield removed, the client stack absorbed the entire 2 MiB body and its terminating chunk before the consumer reached the quarter mark; on an 8 MiB body the server was ~5.8 MiB in at the stall, and on 32 MiB it was ~13.7 MiB in. So ~500 KiB is not a ceiling this 2 MiB body clears — the buffers here hold several MiB. (A "~526 KiB" figure measured against a consumer that has already stopped reading does not bound read-ahead during the phase where the consumer is reading flat out.)

Under CI conditions the property does break. python_sdk_tests.yml runs --numprocesses=4; with that plus CPU oversubscription, 1 of 12 runs here had the stall land after the whole body, terminating chunk included (2048/2048, body_complete=True), and several async runs drifted to 1155–1578. Degenerate runs still pass — they just stop covering slow-consumer behavior. That matters because the wall-clock mutation this test was validated against is also caught by the pre-existing test_{sync,async}_stream_survives_transfers_longer_than_read_timeout. The unique coverage here is the consumer-pace case, which needs a socket read still outstanding after the stall — precisely what the degenerate runs remove.

Verified alternative: pace the server rather than trying to out-size the buffers, and assert the property. The fixture already takes per-chunk delays, so [0.005] * 400 puts a hard 2s floor under the transfer and a 0.9s stall at the quarter mark cannot land past it:

SLOW_CONSUMER_CHUNKS = 400
SLOW_CONSUMER_PACE = 0.005  # every gap far under the idle bound
SLOW_CONSUMER_STALL_AT = SLOW_CONSUMER_CHUNKS // 4

api_url, server = _start_volume_file_server([SLOW_CONSUMER_PACE] * SLOW_CONSUMER_CHUNKS)
...
    if not stalled and received >= len(CHUNK) * SLOW_CONSUMER_STALL_AT:
        stalled = True
        # The stall has to land while the server is still pushing, or the case
        # degenerates into re-reading an already-buffered body.
        assert not server.body_complete
        time.sleep(short_read_timeout * 3)

Measured: 99–100 of 400 chunks at the stall in 10/10 runs across both flavors, all under the same 8-busy-loop load that produced the degenerate run above; the body drops from 2 MiB to 400 KiB. It needs the helper to hand back a small progress object next to the URL — tests/envd_frame_server.py's SharedPoolServer.drop_when is the in-repo precedent for a test server the test synchronizes with rather than races.

Non-blocking on correctness, but the point of the PR is coverage that cannot pass vacuously, and this is the one test here that currently can.

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

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This claim is slightly off, and the correction is in the test's favour rather than against it.

I mutated the streaming path so the default bound becomes a whole-transfer deadline (_get_request_timeout(READ_TIMEOUT, ...) instead of (None, ...), sync and async). Six tests fail, only two of which are new — test_{sync,async}_stream_survives_transfers_longer_than_read_timeout, test_async_explicit_stream_idle_timeout_above_transport_bound and test_async_stream_idle_timeout_zero_disables_idle_bound already catch it. The two genuine stall tests do pass, so the sentence is literally true as written, but a reader will take it to mean this test is the guard against that regression, and it is not.

The regression it uniquely catches is a narrower and more interesting one: an idle bound measured across consumer time rather than wire time. I simulated that by starting the idle timer before the yield instead of after, and exactly one test in the file fails — this one. The five pre-existing timing tests all pass, because none of them ever pauses between reads.

So the comment would be both stronger and accurate as something like: "the bound must stay a wire-only one — a bound that also counted the consumer's own pauses would fail here and nowhere else in this file."

# 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")
Expand All @@ -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(

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nit, but a cheap one to close: this A/B counterpart exists only for sync, while the two flavors reach the deadline through different code — the pyqwest adapter's async path re-applies the remaining deadline around each chunk read (AsyncIteratorByteStream), whereas the sync path relies on reqwest keeping the deadline attached to the in-flight request. So "a regression that turned the default idle bound into a wall-clock one would look indistinguishable from correct behavior" is only ruled out for one of the two mechanisms, while every other streaming case in this file is paired sync/async.

I checked that the async twin would actually pass rather than just suggesting it: same fixture and body, request_timeout=1.0 with a 2.0s stall at the quarter mark raises httpx.ReadTimeout (a httpx.TimeoutException) in 3/3 runs, so it works with the assertion already written here.

Smaller point in the same test: request_timeout=1.0 and time.sleep(2.0) are the two numbers that make it valid — the stall has to exceed the deadline — and unlike the SLOW_CONSUMER_* constants above they are unnamed, so the relationship between them lives only in the prose.

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

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is the most precisely targeted test in the PR, and the A/B framing pays off: dropping the caller's explicit request_timeout on the streaming path (_get_request_timeout(None, None)) fails only this test, with all 9 of its neighbours green. Nothing else in the suite pins that the explicit value reaches the stream at all.

One caveat on the reasoning in the last sentence. "The tests above establish that a slow consumer never trips the idle bound" holds against the current implementation, but the two tests above and this one share the same fixture, body size and stall, so a change to SLOW_CONSUMER_STALL_AT or the sleep durations moves all three together and the A/B silently stops being an A/B. Since the deduction is what gives this test its meaning, it is worth making the relationship mechanical — deriving the 2.0 s stall from short_read_timeout the way the other two do, rather than hard-coding it against a hard-coded request_timeout=1.0.

# 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")
Expand Down
Loading