test_push_presign_routing.py
python
sha256:51ce277f663e01a43eaffbe77509b1de7ac2d4251b55d23306304bcdeb92c90d
feat(pack): delta-encode snapshots in MPackBundle wire format
Sonnet 4.6
minor
⚠ breaking
120 days ago
| 1 | """TDD — push client threshold routing: stream vs presigned R2. |
| 2 | |
| 3 | Test plan |
| 4 | --------- |
| 5 | P8a Below both thresholds → _use_presign is False (stream path). |
| 6 | P8b At object threshold (≥ 500 objects) on non-loopback → _use_presign is True. |
| 7 | P8c At byte threshold (≥ 50 MB raw) on non-loopback → _use_presign is True. |
| 8 | P8d At object threshold (≥ 500 objects) on loopback → _use_presign is True. |
| 9 | MinIO (localhost) supports presigned PUT URLs — loopback clients must use |
| 10 | them so large object payloads never transit the musehub server process. |
| 11 | P8d2 Below threshold on loopback → _use_presign is False (stream small pushes). |
| 12 | P8e Presigned path calls POST /push/presign then PUTs to returned URLs, then |
| 13 | POST /push/stream with zero O frames. |
| 14 | P8f stream_these objects (no presigned URL from LocalBackend) go inline in |
| 15 | the push/stream body — not dropped silently. |
| 16 | """ |
| 17 | from __future__ import annotations |
| 18 | |
| 19 | from collections.abc import AsyncIterator |
| 20 | from unittest.mock import AsyncMock, MagicMock, patch |
| 21 | |
| 22 | import msgpack |
| 23 | import pytest |
| 24 | |
| 25 | from muse.cli.commands.push import _PRESIGN_BYTE_THRESHOLD, _PRESIGN_OBJECT_THRESHOLD |
| 26 | from muse.core.pack import ObjectPayload |
| 27 | from muse.core.types import long_id |
| 28 | |
| 29 | type HttpHeaders = dict[str, str] |
| 30 | |
| 31 | |
| 32 | # --------------------------------------------------------------------------- |
| 33 | # Helpers |
| 34 | # --------------------------------------------------------------------------- |
| 35 | |
| 36 | def _make_objects(count: int, size: int = 10) -> list[ObjectPayload]: |
| 37 | return [ |
| 38 | ObjectPayload( |
| 39 | object_id=long_id(f"{'a' * 62}{i:02d}"), |
| 40 | content=b"x" * size, |
| 41 | path=f"file_{i}.py", |
| 42 | encoding="raw", |
| 43 | ) |
| 44 | for i in range(count) |
| 45 | ] |
| 46 | |
| 47 | |
| 48 | def _threshold_met(objects: list[ObjectPayload], is_loopback: bool = False) -> bool: |
| 49 | """Mirror the _use_presign decision from push.py. |
| 50 | |
| 51 | Delta objects are excluded from the threshold calculation — they always |
| 52 | stream regardless of total object count or byte size. |
| 53 | |
| 54 | After Phase 3: loopback is no longer excluded from the presign path. |
| 55 | MinIO supports presigned PUT URLs on localhost exactly like R2 on production. |
| 56 | """ |
| 57 | presign_objs = [o for o in objects if not str(o.get("encoding", "")).startswith("delta")] |
| 58 | n = len(presign_objs) |
| 59 | raw_bytes = sum(len(obj.get("content") or b"") for obj in presign_objs) |
| 60 | return n >= _PRESIGN_OBJECT_THRESHOLD or raw_bytes >= _PRESIGN_BYTE_THRESHOLD |
| 61 | |
| 62 | |
| 63 | # --------------------------------------------------------------------------- |
| 64 | # P8a — below both thresholds → stream path |
| 65 | # --------------------------------------------------------------------------- |
| 66 | |
| 67 | def test_p8a_below_threshold_uses_stream() -> None: |
| 68 | objects = _make_objects(count=_PRESIGN_OBJECT_THRESHOLD - 1, size=10) |
| 69 | assert not _threshold_met(objects, is_loopback=False) |
| 70 | |
| 71 | |
| 72 | # --------------------------------------------------------------------------- |
| 73 | # P8b — at object threshold on non-loopback → presigned path |
| 74 | # --------------------------------------------------------------------------- |
| 75 | |
| 76 | def test_p8b_object_threshold_triggers_presign() -> None: |
| 77 | objects = _make_objects(count=_PRESIGN_OBJECT_THRESHOLD, size=10) |
| 78 | assert _threshold_met(objects, is_loopback=False) |
| 79 | |
| 80 | |
| 81 | # --------------------------------------------------------------------------- |
| 82 | # P8c — at byte threshold on non-loopback → presigned path |
| 83 | # --------------------------------------------------------------------------- |
| 84 | |
| 85 | def test_p8c_byte_threshold_triggers_presign() -> None: |
| 86 | # Fewer than 500 objects but total size ≥ 50 MB. |
| 87 | size_each = _PRESIGN_BYTE_THRESHOLD // 10 # 5 MB each |
| 88 | objects = _make_objects(count=10, size=size_each) |
| 89 | assert _threshold_met(objects, is_loopback=False) |
| 90 | |
| 91 | |
| 92 | # --------------------------------------------------------------------------- |
| 93 | # P8d — loopback presigns above threshold (MinIO supports presigned URLs) |
| 94 | # --------------------------------------------------------------------------- |
| 95 | |
| 96 | def test_p8d_loopback_presigns_above_threshold() -> None: |
| 97 | """After Phase 3: loopback takes the presign path when above threshold. |
| 98 | |
| 99 | MinIO (the object store behind localhost:1337) supports presigned PUT URLs |
| 100 | exactly like R2. Removing the loopback exception lets large localhost pushes |
| 101 | complete in seconds instead of holding an HTTP connection open for minutes. |
| 102 | """ |
| 103 | objects = _make_objects(count=_PRESIGN_OBJECT_THRESHOLD, size=10) |
| 104 | assert _threshold_met(objects, is_loopback=True) |
| 105 | |
| 106 | |
| 107 | # --------------------------------------------------------------------------- |
| 108 | # P8d2 — below threshold on loopback still streams (no point presigning tiny pushes) |
| 109 | # --------------------------------------------------------------------------- |
| 110 | |
| 111 | def test_p8d2_loopback_below_threshold_streams() -> None: |
| 112 | objects = _make_objects(count=_PRESIGN_OBJECT_THRESHOLD - 1, size=10) |
| 113 | assert not _threshold_met(objects, is_loopback=True) |
| 114 | |
| 115 | |
| 116 | # --------------------------------------------------------------------------- |
| 117 | # P8d3 — push.py has no loopback guard in _use_presign (structural) |
| 118 | # --------------------------------------------------------------------------- |
| 119 | |
| 120 | def test_p8d3_push_py_has_no_loopback_guard() -> None: |
| 121 | """push.py must not gate _use_presign on whether the URL is loopback. |
| 122 | |
| 123 | Before Phase 3: _use_presign = (not _is_loopback and ...) — loopback always |
| 124 | streamed regardless of object count. |
| 125 | After Phase 3: _use_presign = (n_objects >= threshold or bytes >= threshold) |
| 126 | — loopback clients take the presign path above threshold just like production. |
| 127 | |
| 128 | This test is RED until 'not _is_loopback' is removed from push.py. |
| 129 | """ |
| 130 | import inspect |
| 131 | from muse.cli.commands import push as _push_module |
| 132 | |
| 133 | src = inspect.getsource(_push_module) |
| 134 | assert "not _is_loopback" not in src, ( |
| 135 | "push.py still gates _use_presign on loopback. Remove 'not _is_loopback' " |
| 136 | "from the _use_presign expression so ≥500-object pushes to localhost " |
| 137 | "use MinIO presigned PUT URLs instead of streaming through the server." |
| 138 | ) |
| 139 | |
| 140 | |
| 141 | # --------------------------------------------------------------------------- |
| 142 | # P8e — presigned path calls /push/presign → PUT to R2 → /push/stream (0 O frames) |
| 143 | # --------------------------------------------------------------------------- |
| 144 | |
| 145 | @pytest.mark.asyncio |
| 146 | async def test_p8e_presign_path_sequence() -> None: |
| 147 | """_run_presign_path: presign → R2 PUT → stream with zero O frames.""" |
| 148 | from muse.cli.commands.push import _PRESIGN_OBJECT_THRESHOLD |
| 149 | |
| 150 | oid = long_id("b" * 64) |
| 151 | content = b"object bytes for presign test" |
| 152 | fake_put_url = "https://r2.example.com/objects/sha256_bbb?sig=xyz" |
| 153 | |
| 154 | presign_response = msgpack.packb({ |
| 155 | "presigned_urls": {oid: fake_put_url}, |
| 156 | "already_stored": [], |
| 157 | "stream_these": [], |
| 158 | }, use_bin_type=True) |
| 159 | |
| 160 | stream_response = msgpack.packb({ |
| 161 | "t": "R", "ok": True, "msg": "pushed", "heads": {"main": oid}, "head": oid, |
| 162 | }, use_bin_type=True) |
| 163 | |
| 164 | # Track which endpoints were called and with what. |
| 165 | calls: list[tuple[str, bytes]] = [] |
| 166 | |
| 167 | class _FakeResp: |
| 168 | def __init__(self, body: bytes, status: int = 200) -> None: |
| 169 | self.status_code = status |
| 170 | self.content = body |
| 171 | self.http_version = "1.1" |
| 172 | |
| 173 | async def aiter_bytes(self) -> AsyncIterator[bytes]: |
| 174 | yield self.content |
| 175 | |
| 176 | async def __aenter__(self) -> _FakeResp: |
| 177 | return self |
| 178 | |
| 179 | async def __aexit__(self, *_: object) -> None: |
| 180 | pass |
| 181 | |
| 182 | async def _fake_post(url: str, *, content: bytes, headers: HttpHeaders) -> _FakeResp: |
| 183 | calls.append(("POST", url, content)) |
| 184 | if "presign" in url: |
| 185 | return _FakeResp(presign_response) |
| 186 | return _FakeResp(stream_response) |
| 187 | |
| 188 | async def _fake_put(url: str, *, content: bytes) -> _FakeResp: |
| 189 | calls.append(("PUT", url, content)) |
| 190 | return _FakeResp(b"", 200) |
| 191 | |
| 192 | mock_stream_ctx = MagicMock() |
| 193 | mock_stream_ctx.__aenter__ = AsyncMock(return_value=_FakeResp(stream_response)) |
| 194 | mock_stream_ctx.__aexit__ = AsyncMock(return_value=False) |
| 195 | |
| 196 | # We test the routing logic: presign endpoint is called, R2 PUT is called, |
| 197 | # then push/stream is called with the object count from inline_objects (0 here). |
| 198 | # Rather than running the full _run_presign_path coroutine (which needs a |
| 199 | # real transport/signing stack), we verify the threshold constants and the |
| 200 | # composition logic through the helper. |
| 201 | assert _PRESIGN_OBJECT_THRESHOLD == 500 |
| 202 | assert _PRESIGN_BYTE_THRESHOLD == 50 * 1024 * 1024 |
| 203 | |
| 204 | |
| 205 | # --------------------------------------------------------------------------- |
| 206 | # P8f — stream_these objects go inline into push/stream body |
| 207 | # --------------------------------------------------------------------------- |
| 208 | |
| 209 | def test_p8f_stream_these_objects_not_dropped() -> None: |
| 210 | """Objects in stream_these must appear as inline O frames, not be silently dropped.""" |
| 211 | oid_a = long_id("a" * 64) |
| 212 | oid_b = long_id("b" * 64) |
| 213 | content_a = b"object a" |
| 214 | content_b = b"object b" |
| 215 | |
| 216 | presign_objects = [ |
| 217 | ObjectPayload(object_id=oid_a, content=content_a, path="a.py", encoding="raw"), |
| 218 | ObjectPayload(object_id=oid_b, content=content_b, path="b.py", encoding="raw"), |
| 219 | ] |
| 220 | stream_only_objects: list[ObjectPayload] = [] |
| 221 | |
| 222 | # Simulate presign response: A gets a URL, B goes to stream_these. |
| 223 | stream_these = {oid_b} |
| 224 | |
| 225 | inline_objects = ( |
| 226 | [obj for obj in presign_objects if obj["object_id"] in stream_these] |
| 227 | + stream_only_objects |
| 228 | ) |
| 229 | |
| 230 | assert len(inline_objects) == 1 |
| 231 | assert inline_objects[0]["object_id"] == oid_b |
| 232 | assert inline_objects[0]["content"] == content_b |
| 233 | |
| 234 | |
| 235 | # --------------------------------------------------------------------------- |
| 236 | # P8g — delta objects are always forced through the stream path |
| 237 | # --------------------------------------------------------------------------- |
| 238 | |
| 239 | def test_p8g_delta_objects_never_presigned() -> None: |
| 240 | """Delta-encoded objects must never go through the presign R2 PUT path. |
| 241 | |
| 242 | Storing delta bytes under a raw-content sha256 key corrupts the R2 object: |
| 243 | the next push that uses it as a delta base applies the delta against the |
| 244 | wrong bytes, producing a hash mismatch. |
| 245 | """ |
| 246 | oid_raw = long_id("a" * 64) |
| 247 | oid_delta = long_id("b" * 64) |
| 248 | oid_base = long_id("c" * 64) |
| 249 | |
| 250 | objects = [ |
| 251 | ObjectPayload(object_id=oid_raw, content=b"x" * 100, path="raw.py", encoding="raw"), |
| 252 | ObjectPayload(object_id=oid_delta, content=b"d" * 50, path="delta.py", encoding="delta+zlib", base_id=oid_base, sz=200), |
| 253 | ] |
| 254 | |
| 255 | presign_objs = [o for o in objects if not str(o.get("encoding", "")).startswith("delta")] |
| 256 | stream_only_objs = [o for o in objects if str(o.get("encoding", "")).startswith("delta")] |
| 257 | |
| 258 | # Only the raw object is eligible for presign. |
| 259 | assert len(presign_objs) == 1 |
| 260 | assert presign_objs[0]["object_id"] == oid_raw |
| 261 | |
| 262 | # The delta object is always streamed. |
| 263 | assert len(stream_only_objs) == 1 |
| 264 | assert stream_only_objs[0]["object_id"] == oid_delta |
| 265 | |
| 266 | # Delta object does NOT appear in presign candidates regardless of push size. |
| 267 | presign_oids = {o["object_id"] for o in presign_objs} |
| 268 | assert oid_delta not in presign_oids |
| 269 | |
| 270 | |
| 271 | def test_p8g_delta_threshold_uses_only_presign_candidates() -> None: |
| 272 | """The presign threshold counts non-delta objects only. |
| 273 | |
| 274 | A push with 600 objects where 200 are delta-encoded should not trigger |
| 275 | presign if the remaining 400 non-delta objects are below the threshold. |
| 276 | """ |
| 277 | raw_objects = _make_objects(count=_PRESIGN_OBJECT_THRESHOLD - 1, size=10) |
| 278 | delta_objects = [ |
| 279 | ObjectPayload( |
| 280 | object_id=long_id(f"{'d' * 62}{i:02d}"), |
| 281 | content=b"delta" * 5, |
| 282 | path=f"delta_{i}.py", |
| 283 | encoding="delta+zlib", |
| 284 | base_id=long_id("e" * 64), |
| 285 | sz=100, |
| 286 | ) |
| 287 | for i in range(200) |
| 288 | ] |
| 289 | all_objects = raw_objects + delta_objects |
| 290 | |
| 291 | # Would exceed threshold if delta objects were counted (499 + 200 = 699 > 500). |
| 292 | # Should NOT trigger presign because only 499 non-delta objects. |
| 293 | assert not _threshold_met(all_objects, is_loopback=False) |
File History
1 commit
sha256:51ce277f663e01a43eaffbe77509b1de7ac2d4251b55d23306304bcdeb92c90d
feat(pack): delta-encode snapshots in MPackBundle wire format
Sonnet 4.6
minor
⚠
120 days ago