"""TDD — presigned PUT concurrency and keepalive. Test plan --------- C1 PUTs are issued concurrently — N objects take < N × single-PUT-latency (proves asyncio.gather is used, not a sequential for loop). C2 The httpx client is created without max_keepalive_connections=0 so TCP connections are reused across the batch (keepalive enabled). C3 A semaphore caps concurrent connections — no unbounded fan-out with 8000+ objects (max_concurrent ≤ PRESIGN_PUT_CONCURRENCY). C4 All objects are uploaded — no silent drops when concurrency is bounded. C5 An HTTP error on one PUT raises TransportError (error handling survives the gather refactor). """ from __future__ import annotations import asyncio import time from unittest.mock import AsyncMock, MagicMock, patch import pytest from muse.cli.commands.push import ( PRESIGN_PUT_CONCURRENCY, _run_presign_puts, ) from muse.core.types import long_id # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- def _make_url_map(n: int) -> dict[str, str]: return {long_id(f"{'a' * 62}{i:02d}"): f"http://localhost:9000/bucket/obj{i}?sig=x" for i in range(n)} def _make_obj_map(url_map: dict[str, str], size: int = 64) -> dict[str, dict]: return { oid: {"object_id": oid, "content": b"x" * size, "path": f"file_{i}.py", "encoding": "raw"} for i, oid in enumerate(url_map) } # --------------------------------------------------------------------------- # C1 — PUTs are concurrent (gather, not sequential for loop) # --------------------------------------------------------------------------- async def test_c1_puts_are_concurrent() -> None: """With N objects each taking DELAY seconds, total time must be well under N × DELAY — proving concurrent dispatch, not sequential.""" N = 20 DELAY = 0.04 # 40 ms per PUT url_map = _make_url_map(N) obj_map = _make_obj_map(url_map) async def _slow_put(url: str, *, content: bytes) -> MagicMock: await asyncio.sleep(DELAY) resp = MagicMock() resp.status_code = 200 return resp mock_client = MagicMock() mock_client.put = _slow_put start = time.perf_counter() await _run_presign_puts(mock_client, url_map, obj_map) elapsed = time.perf_counter() - start # Sequential would take N × DELAY = 0.8s; concurrent must finish in < 2× DELAY. assert elapsed < DELAY * 3, ( f"PUTs appear sequential: {elapsed:.3f}s for {N} objects at {DELAY}s each. " "Expected concurrent execution via asyncio.gather." ) # --------------------------------------------------------------------------- # C2 — keepalive is not disabled on the httpx client # --------------------------------------------------------------------------- def test_c2_keepalive_not_disabled() -> None: """The httpx.Limits passed to AsyncClient must not set max_keepalive_connections=0.""" import httpx captured_limits: list[httpx.Limits] = [] original_init = httpx.AsyncClient.__init__ def _patched_init(self, *args, **kwargs): # type: ignore[no-untyped-def] if "limits" in kwargs: captured_limits.append(kwargs["limits"]) original_init(self, *args, **kwargs) with patch.object(httpx.AsyncClient, "__init__", _patched_init): # Import the module-level client factory used by _run_presign_puts from muse.cli.commands.push import _make_r2_client client_ctx = _make_r2_client() assert captured_limits, "httpx.AsyncClient was not constructed with explicit Limits" for limits in captured_limits: assert limits.max_keepalive_connections != 0, ( "max_keepalive_connections=0 disables keepalive — must be None or a positive int" ) # --------------------------------------------------------------------------- # C3 — semaphore caps concurrent connections # --------------------------------------------------------------------------- async def test_c3_semaphore_caps_concurrency() -> None: """At most PRESIGN_PUT_CONCURRENCY PUTs must be in flight simultaneously.""" N = PRESIGN_PUT_CONCURRENCY * 3 url_map = _make_url_map(N) obj_map = _make_obj_map(url_map) in_flight: list[int] = [] peak: list[int] = [0] current = [0] async def _counting_put(url: str, *, content: bytes) -> MagicMock: current[0] += 1 peak[0] = max(peak[0], current[0]) in_flight.append(current[0]) await asyncio.sleep(0.005) current[0] -= 1 resp = MagicMock() resp.status_code = 200 return resp mock_client = MagicMock() mock_client.put = _counting_put await _run_presign_puts(mock_client, url_map, obj_map) assert peak[0] <= PRESIGN_PUT_CONCURRENCY, ( f"Peak concurrency {peak[0]} exceeded semaphore limit {PRESIGN_PUT_CONCURRENCY}" ) assert peak[0] > 1, "Expected concurrent PUTs but only 1 was ever in flight" # --------------------------------------------------------------------------- # C4 — all objects uploaded despite bounded concurrency # --------------------------------------------------------------------------- async def test_c4_all_objects_uploaded() -> None: """Every object_id in url_map must be PUT exactly once.""" N = PRESIGN_PUT_CONCURRENCY * 2 + 7 # intentionally not a multiple url_map = _make_url_map(N) obj_map = _make_obj_map(url_map) uploaded: list[str] = [] async def _recording_put(url: str, *, content: bytes) -> MagicMock: uploaded.append(url) resp = MagicMock() resp.status_code = 200 return resp mock_client = MagicMock() mock_client.put = _recording_put await _run_presign_puts(mock_client, url_map, obj_map) assert len(uploaded) == N, f"Expected {N} PUTs, got {len(uploaded)}" assert len(set(uploaded)) == N, "Duplicate PUTs detected" # --------------------------------------------------------------------------- # C5 — HTTP error propagates correctly from concurrent gather # --------------------------------------------------------------------------- async def test_c5_http_error_raises_transport_error() -> None: """A 403 response from any PUT must raise TransportError, not be swallowed.""" from muse.core.transport import TransportError url_map = _make_url_map(5) obj_map = _make_obj_map(url_map) call_count = [0] async def _failing_put(url: str, *, content: bytes) -> MagicMock: call_count[0] += 1 resp = MagicMock() resp.status_code = 403 if call_count[0] == 3 else 200 return resp mock_client = MagicMock() mock_client.put = _failing_put with pytest.raises((TransportError, Exception)) as exc_info: await _run_presign_puts(mock_client, url_map, obj_map) assert "403" in str(exc_info.value) or "R2" in str(exc_info.value) or exc_info.type.__name__ in ("TransportError", "RuntimeError")