chore: import upstream snapshot with attribution

2026-07-13 12:47:58 +08:00
commit b16403ea71
789 changed files with 115226 additions and 0 deletions
@@ -0,0 +1,226 @@
+"""Unit tests for the streamed-read helpers.
+
+These exercise connection lifecycle (consume / context manager / explicit
+close / idle timeout / abandonment) without hitting a real sandbox, using a
+local chunked HTTP server.
+"""
+
+import asyncio
+import socket
+import threading
+import time
+from typing import Optional
+
+import httpx
+import pytest
+
+from e2b.sandbox.filesystem.filesystem import (
+    AsyncFileStreamReader,
+    FileStreamReader,
+)
+
+CHUNKS = [f"chunk{i}".encode() for i in range(5)]
+EXPECTED = b"".join(CHUNKS)
+
+
+def _start_chunked_server(
+    stall_before: Optional[int] = None,
+    stall_seconds: float = 0.0,
+) -> int:
+    """Start a one-shot HTTP server that replies with a chunked body.
+
+    When ``stall_before`` is set, 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.
+    """
+    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"
+            )
+            for idx, chunk in enumerate(CHUNKS):
+                if stall_before is not None and 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 _active_connections(client) -> int:
+    # Count connections that are still checked out (a leaked/in-use stream),
+    # not the total pool size. A fully consumed stream returns its connection
+    # to the pool, where it may linger as an idle keep-alive entry until the
+    # server-side close is observed; that lingering idle connection is not a
+    # leak. Asserting on total pool size makes this racy under load (the basis
+    # of a CI flake); counting only non-idle connections is deterministic.
+    return sum(1 for conn in client._transport._pool.connections if not conn.is_idle())
+
+
+def _open_stream(client, port, read_timeout: Optional[float] = None):
+    request = client.build_request(
+        "GET", f"http://127.0.0.1:{port}/files", timeout=httpx.Timeout(5.0)
+    )
+    if read_timeout is not None:
+        # Mirror the SDK: the per-chunk `read` timeout bounds idle gaps.
+        request.extensions["timeout"]["read"] = read_timeout
+    return client.send(request, stream=True)
+
+
+def test_sync_full_consume_releases_connection():
+    with httpx.Client() as client:
+        port = _start_chunked_server()
+        reader = FileStreamReader(_open_stream(client, port))
+        assert b"".join(reader) == EXPECTED
+        assert _active_connections(client) == 0
+
+
+def test_sync_context_manager_releases_on_exit():
+    with httpx.Client() as client:
+        port = _start_chunked_server()
+        with FileStreamReader(_open_stream(client, port)) as reader:
+            assert next(iter(reader)) == CHUNKS[0]
+        # Exiting the context releases the connection even though the stream
+        # was only partially consumed.
+        assert _active_connections(client) == 0
+
+
+def test_sync_close_is_idempotent():
+    with httpx.Client() as client:
+        port = _start_chunked_server()
+        reader = FileStreamReader(_open_stream(client, port))
+        reader.close()
+        reader.close()
+        assert _active_connections(client) == 0
+
+
+def test_sync_idle_timeout_releases_connection():
+    with httpx.Client() as client:
+        # The server stalls before the second chunk for longer than the
+        # reader's idle (read) timeout.
+        port = _start_chunked_server(stall_before=1, stall_seconds=0.5)
+        reader = FileStreamReader(_open_stream(client, port, read_timeout=0.05))
+        it = iter(reader)
+        assert next(it)
+        # The stalled read trips the idle timeout, which propagates and
+        # releases the connection.
+        with pytest.raises(httpx.ReadTimeout):
+            next(it)
+        assert _active_connections(client) == 0
+
+
+def test_sync_slow_consumer_does_not_trip_idle_timeout():
+    # The server sends every chunk promptly; the consumer then pauses far
+    # longer than the read (idle) timeout between iterations. httpx's read
+    # timeout only counts while it's waiting on the wire, so a slow consumer
+    # must not trip it (parity with the JS wire-only idle timeout).
+    with httpx.Client() as client:
+        port = _start_chunked_server()
+        reader = FileStreamReader(_open_stream(client, port, read_timeout=0.05))
+        out = []
+        for chunk in reader:
+            out.append(chunk)
+            time.sleep(0.2)
+        assert b"".join(out) == EXPECTED
+
+
+def test_sync_abandoned_reader_is_reclaimed_on_client_close():
+    client = httpx.Client()
+    port = _start_chunked_server()
+    reader = FileStreamReader(_open_stream(client, port))
+    assert _active_connections(client) == 1
+
+    # The sync reader has no GC safety net: dropping it without closing keeps
+    # the connection checked out (an idle timeout would reclaim a stalled one).
+    del reader
+    assert _active_connections(client) == 1
+
+    # Closing the client reclaims the abandoned connection.
+    client.close()
+    assert _active_connections(client) == 0
+
+
+async def test_async_full_consume_releases_connection():
+    async with httpx.AsyncClient() as client:
+        port = _start_chunked_server()
+        request = client.build_request("GET", f"http://127.0.0.1:{port}/files")
+        reader = AsyncFileStreamReader(await client.send(request, stream=True))
+        collected = b"".join([chunk async for chunk in reader])
+        assert collected == EXPECTED
+        assert _active_connections(client) == 0
+
+
+async def test_async_context_manager_releases_on_exit():
+    async with httpx.AsyncClient() as client:
+        port = _start_chunked_server()
+        request = client.build_request("GET", f"http://127.0.0.1:{port}/files")
+        async with AsyncFileStreamReader(
+            await client.send(request, stream=True)
+        ) as reader:
+            assert await reader.__anext__() == CHUNKS[0]
+        assert _active_connections(client) == 0
+
+
+async def test_async_aclose_is_idempotent():
+    async with httpx.AsyncClient() as client:
+        port = _start_chunked_server()
+        request = client.build_request("GET", f"http://127.0.0.1:{port}/files")
+        reader = AsyncFileStreamReader(await client.send(request, stream=True))
+        await reader.aclose()
+        await reader.aclose()
+        assert _active_connections(client) == 0
+
+
+async def test_async_idle_timeout_releases_connection():
+    async with httpx.AsyncClient() as client:
+        port = _start_chunked_server(stall_before=1, stall_seconds=0.5)
+        request = client.build_request(
+            "GET", f"http://127.0.0.1:{port}/files", timeout=httpx.Timeout(5.0)
+        )
+        request.extensions["timeout"]["read"] = 0.05
+        reader = AsyncFileStreamReader(await client.send(request, stream=True))
+        it = reader.__aiter__()
+        assert await it.__anext__()
+        # The stalled read trips the idle timeout, which propagates and
+        # releases the connection.
+        with pytest.raises(httpx.ReadTimeout):
+            await it.__anext__()
+        assert _active_connections(client) == 0
+
+
+async def test_async_abandoned_reader_is_reclaimed_on_client_close():
+    client = httpx.AsyncClient()
+    port = _start_chunked_server()
+    request = client.build_request("GET", f"http://127.0.0.1:{port}/files")
+    reader = AsyncFileStreamReader(await client.send(request, stream=True))
+    assert _active_connections(client) == 1
+
+    # The async reader has no GC safety net: dropping it without closing keeps
+    # the connection checked out (releasing one requires awaiting aclose()).
+    del reader
+    await asyncio.sleep(0.05)
+    assert _active_connections(client) == 1
+
+    # Closing the client reclaims the abandoned connection.
+    await client.aclose()
+    assert _active_connections(client) == 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(pytest.main([__file__, "-v"]))