gabriel / muse public
test_push_presign_concurrency.py python
193 lines 6.9 KB
Raw
sha256:51ce277f663e01a43eaffbe77509b1de7ac2d4251b55d23306304bcdeb92c90d feat(pack): delta-encode snapshots in MPackBundle wire format Sonnet 4.6 minor ⚠ breaking 123 days ago
1 """TDD — presigned PUT concurrency and keepalive.
2
3 Test plan
4 ---------
5 C1 PUTs are issued concurrently — N objects take < N × single-PUT-latency
6 (proves asyncio.gather is used, not a sequential for loop).
7 C2 The httpx client is created without max_keepalive_connections=0 so TCP
8 connections are reused across the batch (keepalive enabled).
9 C3 A semaphore caps concurrent connections — no unbounded fan-out with 8000+
10 objects (max_concurrent ≤ PRESIGN_PUT_CONCURRENCY).
11 C4 All objects are uploaded — no silent drops when concurrency is bounded.
12 C5 An HTTP error on one PUT raises TransportError (error handling survives
13 the gather refactor).
14 """
15 from __future__ import annotations
16
17 import asyncio
18 import time
19 from unittest.mock import AsyncMock, MagicMock, patch
20
21 import pytest
22
23 from muse.cli.commands.push import (
24 PRESIGN_PUT_CONCURRENCY,
25 _run_presign_puts,
26 )
27 from muse.core.types import long_id
28
29
30 # ---------------------------------------------------------------------------
31 # Helpers
32 # ---------------------------------------------------------------------------
33
34 def _make_url_map(n: int) -> dict[str, str]:
35 return {long_id(f"{'a' * 62}{i:02d}"): f"http://localhost:9000/bucket/obj{i}?sig=x" for i in range(n)}
36
37
38 def _make_obj_map(url_map: dict[str, str], size: int = 64) -> dict[str, dict]:
39 return {
40 oid: {"object_id": oid, "content": b"x" * size, "path": f"file_{i}.py", "encoding": "raw"}
41 for i, oid in enumerate(url_map)
42 }
43
44
45 # ---------------------------------------------------------------------------
46 # C1 — PUTs are concurrent (gather, not sequential for loop)
47 # ---------------------------------------------------------------------------
48
49 async def test_c1_puts_are_concurrent() -> None:
50 """With N objects each taking DELAY seconds, total time must be well under
51 N × DELAY — proving concurrent dispatch, not sequential."""
52 N = 20
53 DELAY = 0.04 # 40 ms per PUT
54
55 url_map = _make_url_map(N)
56 obj_map = _make_obj_map(url_map)
57
58 async def _slow_put(url: str, *, content: bytes) -> MagicMock:
59 await asyncio.sleep(DELAY)
60 resp = MagicMock()
61 resp.status_code = 200
62 return resp
63
64 mock_client = MagicMock()
65 mock_client.put = _slow_put
66
67 start = time.perf_counter()
68 await _run_presign_puts(mock_client, url_map, obj_map)
69 elapsed = time.perf_counter() - start
70
71 # Sequential would take N × DELAY = 0.8s; concurrent must finish in < 2× DELAY.
72 assert elapsed < DELAY * 3, (
73 f"PUTs appear sequential: {elapsed:.3f}s for {N} objects at {DELAY}s each. "
74 "Expected concurrent execution via asyncio.gather."
75 )
76
77
78 # ---------------------------------------------------------------------------
79 # C2 — keepalive is not disabled on the httpx client
80 # ---------------------------------------------------------------------------
81
82 def test_c2_keepalive_not_disabled() -> None:
83 """The httpx.Limits passed to AsyncClient must not set max_keepalive_connections=0."""
84 import httpx
85 captured_limits: list[httpx.Limits] = []
86
87 original_init = httpx.AsyncClient.__init__
88
89 def _patched_init(self, *args, **kwargs): # type: ignore[no-untyped-def]
90 if "limits" in kwargs:
91 captured_limits.append(kwargs["limits"])
92 original_init(self, *args, **kwargs)
93
94 with patch.object(httpx.AsyncClient, "__init__", _patched_init):
95 # Import the module-level client factory used by _run_presign_puts
96 from muse.cli.commands.push import _make_r2_client
97 client_ctx = _make_r2_client()
98
99 assert captured_limits, "httpx.AsyncClient was not constructed with explicit Limits"
100 for limits in captured_limits:
101 assert limits.max_keepalive_connections != 0, (
102 "max_keepalive_connections=0 disables keepalive — must be None or a positive int"
103 )
104
105
106 # ---------------------------------------------------------------------------
107 # C3 — semaphore caps concurrent connections
108 # ---------------------------------------------------------------------------
109
110 async def test_c3_semaphore_caps_concurrency() -> None:
111 """At most PRESIGN_PUT_CONCURRENCY PUTs must be in flight simultaneously."""
112 N = PRESIGN_PUT_CONCURRENCY * 3
113 url_map = _make_url_map(N)
114 obj_map = _make_obj_map(url_map)
115
116 in_flight: list[int] = []
117 peak: list[int] = [0]
118 current = [0]
119
120 async def _counting_put(url: str, *, content: bytes) -> MagicMock:
121 current[0] += 1
122 peak[0] = max(peak[0], current[0])
123 in_flight.append(current[0])
124 await asyncio.sleep(0.005)
125 current[0] -= 1
126 resp = MagicMock()
127 resp.status_code = 200
128 return resp
129
130 mock_client = MagicMock()
131 mock_client.put = _counting_put
132
133 await _run_presign_puts(mock_client, url_map, obj_map)
134
135 assert peak[0] <= PRESIGN_PUT_CONCURRENCY, (
136 f"Peak concurrency {peak[0]} exceeded semaphore limit {PRESIGN_PUT_CONCURRENCY}"
137 )
138 assert peak[0] > 1, "Expected concurrent PUTs but only 1 was ever in flight"
139
140
141 # ---------------------------------------------------------------------------
142 # C4 — all objects uploaded despite bounded concurrency
143 # ---------------------------------------------------------------------------
144
145 async def test_c4_all_objects_uploaded() -> None:
146 """Every object_id in url_map must be PUT exactly once."""
147 N = PRESIGN_PUT_CONCURRENCY * 2 + 7 # intentionally not a multiple
148 url_map = _make_url_map(N)
149 obj_map = _make_obj_map(url_map)
150
151 uploaded: list[str] = []
152
153 async def _recording_put(url: str, *, content: bytes) -> MagicMock:
154 uploaded.append(url)
155 resp = MagicMock()
156 resp.status_code = 200
157 return resp
158
159 mock_client = MagicMock()
160 mock_client.put = _recording_put
161
162 await _run_presign_puts(mock_client, url_map, obj_map)
163
164 assert len(uploaded) == N, f"Expected {N} PUTs, got {len(uploaded)}"
165 assert len(set(uploaded)) == N, "Duplicate PUTs detected"
166
167
168 # ---------------------------------------------------------------------------
169 # C5 — HTTP error propagates correctly from concurrent gather
170 # ---------------------------------------------------------------------------
171
172 async def test_c5_http_error_raises_transport_error() -> None:
173 """A 403 response from any PUT must raise TransportError, not be swallowed."""
174 from muse.core.transport import TransportError
175
176 url_map = _make_url_map(5)
177 obj_map = _make_obj_map(url_map)
178
179 call_count = [0]
180
181 async def _failing_put(url: str, *, content: bytes) -> MagicMock:
182 call_count[0] += 1
183 resp = MagicMock()
184 resp.status_code = 403 if call_count[0] == 3 else 200
185 return resp
186
187 mock_client = MagicMock()
188 mock_client.put = _failing_put
189
190 with pytest.raises((TransportError, Exception)) as exc_info:
191 await _run_presign_puts(mock_client, url_map, obj_map)
192
193 assert "403" in str(exc_info.value) or "R2" in str(exc_info.value) or exc_info.type.__name__ in ("TransportError", "RuntimeError")
File History 1 commit
sha256:51ce277f663e01a43eaffbe77509b1de7ac2d4251b55d23306304bcdeb92c90d feat(pack): delta-encode snapshots in MPackBundle wire format Sonnet 4.6 minor ⚠ 123 days ago