gabriel / musehub public
test_wire_batch_stream.py python
444 lines 16.1 KB
Raw
sha256:a34090cc4a394a78bd72cbbe34b08cc59525141e19135b6c0ab154f10611b9ef debug(push/stream): instrument O-frame decode path with INF… Sonnet 4.6 patch 121 days ago
1 """TDD — batch streaming: wire_fetch_stream must not hold all snapshots in memory at once.
2
3 The muse wire protocol mirrors git's smart HTTP pack protocol. Git streams a
4 packfile incrementally — O(window_size) memory, not O(full_repo). We must do
5 the same: process commits in batches, emitting O frames for new objects then a
6 C frame for those commits, before moving to the next batch.
7
8 Test plan
9 ---------
10 B1 O frames always precede the C frame that references them.
11 For every commit in a C frame, every object in its snapshot manifest must
12 have appeared in an earlier O frame.
13
14 B2 Each object_id appears in at most one O frame across the whole stream.
15 No redundant downloads.
16
17 B3 A repo whose commit count exceeds _COMMIT_BATCH still delivers all commits.
18 The stream completes correctly when multiple C frames are emitted.
19
20 B4 All objects from all commits are present exactly once across all O frames.
21 No objects missing, no duplicates.
22
23 B5 Commits in different batches that share a snapshot object send that object
24 only once (in the first batch's O frames).
25
26 B6 Empty repo (no commits to fetch) produces only H + E frames — no crash.
27
28 B7 Repo with a single commit spanning >_COMMIT_BATCH objects streams correctly.
29 """
30 from __future__ import annotations
31
32 import asyncio
33 import zlib
34 from datetime import datetime, timezone
35 from typing import Any
36
37 import msgpack
38 import pytest
39 from sqlalchemy.ext.asyncio import AsyncSession
40
41 from muse.core.types import blob_id, fake_id
42 from musehub.db import musehub_models as db
43 from musehub.models.wire import (
44 SFRAME_COMMIT_PACK,
45 SFRAME_END,
46 SFRAME_HEADER,
47 SFRAME_OBJECT,
48 WireFetchRequest,
49 )
50 from musehub.types.json_types import JSONObject, StrDict
51 from tests.factories import create_repo
52
53 type _OidsBySnapshot = dict[str, set[str]]
54
55
56 # ---------------------------------------------------------------------------
57 # Helpers
58 # ---------------------------------------------------------------------------
59
60 def _uid(seed: str) -> str:
61 return fake_id(seed)
62
63
64 def _now() -> datetime:
65 return datetime.now(tz=timezone.utc)
66
67
68 async def _make_objects(
69 session: AsyncSession,
70 repo_id: str,
71 count: int,
72 prefix: str = "obj",
73 ) -> StrDict:
74 """Create *count* objects and return {path: oid} manifest fragment."""
75 from musehub.services.musehub_wire import get_backend
76 from sqlalchemy.dialects.postgresql import insert as pg_insert
77
78 backend = get_backend()
79 manifest: dict[str, str] = {}
80 for i in range(count):
81 content = f"{prefix}-{i}-content".encode()
82 oid = blob_id(content)
83 uri = await backend.put(oid, content)
84 await session.execute(
85 pg_insert(db.MusehubObject)
86 .values(
87 object_id=oid,
88 path=f"{prefix}/{i}.dat",
89 size_bytes=len(content),
90 disk_path=uri.replace("local://", ""),
91 storage_uri=uri,
92 )
93 .on_conflict_do_nothing(index_elements=["object_id"])
94 )
95 await session.execute(
96 pg_insert(db.MusehubObjectRef)
97 .values(repo_id=repo_id, object_id=oid)
98 .on_conflict_do_nothing()
99 )
100 manifest[f"{prefix}/{i}.dat"] = oid
101 await session.commit()
102 return manifest
103
104
105 async def _make_snapshot(
106 session: AsyncSession,
107 repo_id: str,
108 manifest: dict[str, str],
109 ) -> db.MusehubSnapshot:
110 sid = _uid(str(sorted(manifest.items())))
111 snap = db.MusehubSnapshot(
112 snapshot_id=sid,
113 repo_id=repo_id,
114 directories=[],
115 manifest_blob=msgpack.packb(manifest, use_bin_type=True),
116 entry_count=len(manifest),
117 created_at=_now(),
118 )
119 session.add(snap)
120 await session.commit()
121 return snap
122
123
124 async def _make_commit(
125 session: AsyncSession,
126 repo_id: str,
127 *,
128 parent_ids: list[str] | None = None,
129 snapshot_id: str | None = None,
130 seed: str = "",
131 ) -> db.MusehubCommit:
132 commit_id = _uid(f"commit-{seed}")
133 row = db.MusehubCommit(
134 commit_id=commit_id,
135 repo_id=repo_id,
136 branch="main",
137 parent_ids=parent_ids or [],
138 message=f"commit {seed}",
139 author="gabriel",
140 timestamp=_now(),
141 snapshot_id=snapshot_id,
142 )
143 session.add(row)
144 await session.commit()
145 return row
146
147
148 def _unpack_stream(chunks: list[bytes]) -> list[dict]:
149 unpacker = msgpack.Unpacker(raw=False)
150 for chunk in chunks:
151 unpacker.feed(chunk)
152 return list(unpacker)
153
154
155 async def _collect_frames(gen: AsyncIterator[bytes]) -> list[JSONObject]:
156 chunks: list[bytes] = []
157 async for chunk in gen:
158 chunks.append(chunk)
159 return _unpack_stream(chunks)
160
161
162 def _o_frame_oids(frames: list[dict]) -> list[str]:
163 """All object IDs from O frames in order."""
164 return [f["id"] for f in frames if f.get("t") == SFRAME_OBJECT]
165
166
167 def _c_frame_commits(frames: list[dict]) -> list[tuple[int, dict]]:
168 """List of (frame_index, commit_dict) for every commit in every C frame."""
169 result = []
170 for i, f in enumerate(frames):
171 if f.get("t") == SFRAME_COMMIT_PACK:
172 for c in f.get("commits", []):
173 result.append((i, c))
174 return result
175
176
177 def _snapshot_oids(frames: list[JSONObject]) -> _OidsBySnapshot:
178 """Return {snapshot_id: set(oid)} from snapshot data in C frames."""
179 result: dict[str, set[str]] = {}
180 for f in frames:
181 if f.get("t") == SFRAME_COMMIT_PACK:
182 for s in f.get("snapshots", []):
183 sid = s.get("snapshot_id", "")
184 oids = set(s.get("manifest", {}).values())
185 result[sid] = oids
186 return result
187
188
189 # ---------------------------------------------------------------------------
190 # B1 — O frames precede the C frame that references them
191 # ---------------------------------------------------------------------------
192
193 class TestObjectsPrecedeCommits:
194 @pytest.mark.asyncio
195 async def test_objects_arrive_before_their_commit(
196 self, db_session: AsyncSession
197 ) -> None:
198 """Every object in a commit's snapshot must appear in an O frame
199 that comes before the C frame containing that commit."""
200 repo = await create_repo(db_session, owner="test-batch-b1")
201 repo_id = str(repo.repo_id)
202
203 manifest = await _make_objects(db_session, repo_id, 5, prefix="b1")
204 snap = await _make_snapshot(db_session, repo_id, manifest)
205 c1 = await _make_commit(db_session, repo_id, snapshot_id=snap.snapshot_id, seed="b1")
206
207 from musehub.services.musehub_wire import wire_fetch_stream
208 req = WireFetchRequest(want=[c1.commit_id], have=[])
209 frames = await _collect_frames(wire_fetch_stream(db_session, repo_id, req))
210
211 seen_oids: set[str] = set()
212 snap_oids = _snapshot_oids(frames)
213
214 for i, frame in enumerate(frames):
215 if frame.get("t") == SFRAME_OBJECT:
216 seen_oids.add(frame["id"])
217 elif frame.get("t") == SFRAME_COMMIT_PACK:
218 for commit in frame.get("commits", []):
219 sid = commit.get("snapshot_id", "")
220 needed = snap_oids.get(sid, set())
221 missing = needed - seen_oids
222 assert not missing, (
223 f"C frame at index {i} references objects not yet sent: {missing}"
224 )
225
226
227 # ---------------------------------------------------------------------------
228 # B2 — Each object appears in at most one O frame
229 # ---------------------------------------------------------------------------
230
231 class TestNoDuplicateObjects:
232 @pytest.mark.asyncio
233 async def test_no_object_sent_twice(self, db_session: AsyncSession) -> None:
234 repo = await create_repo(db_session, owner="test-batch-b2")
235 repo_id = str(repo.repo_id)
236
237 manifest = await _make_objects(db_session, repo_id, 10, prefix="b2")
238 snap = await _make_snapshot(db_session, repo_id, manifest)
239 c1 = await _make_commit(db_session, repo_id, snapshot_id=snap.snapshot_id, seed="b2a")
240 c2 = await _make_commit(
241 db_session, repo_id,
242 parent_ids=[c1.commit_id],
243 snapshot_id=snap.snapshot_id, # same snapshot — objects must not repeat
244 seed="b2b",
245 )
246
247 from musehub.services.musehub_wire import wire_fetch_stream
248 req = WireFetchRequest(want=[c2.commit_id], have=[])
249 frames = await _collect_frames(wire_fetch_stream(db_session, repo_id, req))
250
251 oids = _o_frame_oids(frames)
252 assert len(oids) == len(set(oids)), (
253 f"Duplicate object IDs in stream: {[o for o in oids if oids.count(o) > 1]}"
254 )
255
256
257 # ---------------------------------------------------------------------------
258 # B3 — Repo with > _COMMIT_BATCH commits streams all commits across multiple C frames
259 # ---------------------------------------------------------------------------
260
261 class TestMultiBatchCompletion:
262 @pytest.mark.asyncio
263 async def test_all_commits_delivered_across_batches(
264 self, db_session: AsyncSession
265 ) -> None:
266 """A chain of 60 commits (> _COMMIT_BATCH=50) must all appear in the stream."""
267 from musehub.services.musehub_wire import _COMMIT_BATCH
268
269 repo = await create_repo(db_session, owner="test-batch-b3")
270 repo_id = str(repo.repo_id)
271
272 n = _COMMIT_BATCH + 10 # one more batch than the limit
273 manifest = await _make_objects(db_session, repo_id, 3, prefix="b3")
274 snap = await _make_snapshot(db_session, repo_id, manifest)
275
276 prev_id: str | None = None
277 commit_ids: list[str] = []
278 for i in range(n):
279 c = await _make_commit(
280 db_session, repo_id,
281 parent_ids=[prev_id] if prev_id else [],
282 snapshot_id=snap.snapshot_id,
283 seed=f"b3-{i}",
284 )
285 prev_id = c.commit_id
286 commit_ids.append(c.commit_id)
287
288 from musehub.services.musehub_wire import wire_fetch_stream
289 req = WireFetchRequest(want=[prev_id], have=[])
290 frames = await _collect_frames(wire_fetch_stream(db_session, repo_id, req))
291
292 c_frames = [f for f in frames if f.get("t") == SFRAME_COMMIT_PACK]
293 assert len(c_frames) >= 2, "Expected multiple C frames for a large repo"
294
295 received_ids = {
296 c["commit_id"]
297 for cf in c_frames
298 for c in cf.get("commits", [])
299 }
300 missing = set(commit_ids) - received_ids
301 assert not missing, f"Missing commits after batch stream: {missing}"
302
303
304 # ---------------------------------------------------------------------------
305 # B4 — All objects from all commits present exactly once
306 # ---------------------------------------------------------------------------
307
308 class TestAllObjectsPresent:
309 @pytest.mark.asyncio
310 async def test_all_unique_objects_sent_exactly_once(
311 self, db_session: AsyncSession
312 ) -> None:
313 repo = await create_repo(db_session, owner="test-batch-b4")
314 repo_id = str(repo.repo_id)
315
316 # Three commits, each adding distinct objects
317 m1 = await _make_objects(db_session, repo_id, 4, prefix="b4-snap1")
318 snap1 = await _make_snapshot(db_session, repo_id, m1)
319 c1 = await _make_commit(db_session, repo_id, snapshot_id=snap1.snapshot_id, seed="b4-1")
320
321 m2 = {**m1, **await _make_objects(db_session, repo_id, 4, prefix="b4-snap2")}
322 snap2 = await _make_snapshot(db_session, repo_id, m2)
323 c2 = await _make_commit(
324 db_session, repo_id,
325 parent_ids=[c1.commit_id],
326 snapshot_id=snap2.snapshot_id,
327 seed="b4-2",
328 )
329
330 m3 = {**m2, **await _make_objects(db_session, repo_id, 4, prefix="b4-snap3")}
331 snap3 = await _make_snapshot(db_session, repo_id, m3)
332 c3 = await _make_commit(
333 db_session, repo_id,
334 parent_ids=[c2.commit_id],
335 snapshot_id=snap3.snapshot_id,
336 seed="b4-3",
337 )
338
339 all_oids = set(m3.values()) # union of all three manifests
340
341 from musehub.services.musehub_wire import wire_fetch_stream
342 req = WireFetchRequest(want=[c3.commit_id], have=[])
343 frames = await _collect_frames(wire_fetch_stream(db_session, repo_id, req))
344
345 received = _o_frame_oids(frames)
346 received_set = set(received)
347
348 assert received_set == all_oids, (
349 f"Missing objects: {all_oids - received_set}\n"
350 f"Extra objects: {received_set - all_oids}"
351 )
352 assert len(received) == len(received_set), "Duplicate objects in stream"
353
354
355 # ---------------------------------------------------------------------------
356 # B5 — Shared objects across commit batches are sent only once
357 # ---------------------------------------------------------------------------
358
359 class TestSharedObjectsSentOnce:
360 @pytest.mark.asyncio
361 async def test_shared_snapshot_objects_not_duplicated(
362 self, db_session: AsyncSession
363 ) -> None:
364 """Two commits in separate batches sharing the same snapshot objects
365 must not cause those objects to be sent twice."""
366 from musehub.services.musehub_wire import _COMMIT_BATCH
367
368 repo = await create_repo(db_session, owner="test-batch-b5")
369 repo_id = str(repo.repo_id)
370
371 shared = await _make_objects(db_session, repo_id, 5, prefix="b5-shared")
372 snap_shared = await _make_snapshot(db_session, repo_id, shared)
373
374 # Build a chain long enough to span two batches
375 prev_id: str | None = None
376 for i in range(_COMMIT_BATCH + 1):
377 c = await _make_commit(
378 db_session, repo_id,
379 parent_ids=[prev_id] if prev_id else [],
380 snapshot_id=snap_shared.snapshot_id,
381 seed=f"b5-{i}",
382 )
383 prev_id = c.commit_id
384
385 from musehub.services.musehub_wire import wire_fetch_stream
386 req = WireFetchRequest(want=[prev_id], have=[])
387 frames = await _collect_frames(wire_fetch_stream(db_session, repo_id, req))
388
389 oids = _o_frame_oids(frames)
390 assert len(oids) == len(set(oids)), "Shared objects sent more than once"
391 # Exactly the shared objects, no more
392 assert set(oids) == set(shared.values())
393
394
395 # ---------------------------------------------------------------------------
396 # B6 — Empty fetch (nothing to send) produces H + E only
397 # ---------------------------------------------------------------------------
398
399 class TestEmptyFetch:
400 @pytest.mark.asyncio
401 async def test_empty_want_produces_header_and_end(
402 self, db_session: AsyncSession
403 ) -> None:
404 repo = await create_repo(db_session, owner="test-batch-b6")
405 repo_id = str(repo.repo_id)
406
407 from musehub.services.musehub_wire import wire_fetch_stream
408 req = WireFetchRequest(want=[], have=[])
409 frames = await _collect_frames(wire_fetch_stream(db_session, repo_id, req))
410
411 types = [f.get("t") for f in frames]
412 assert SFRAME_HEADER in types
413 assert SFRAME_END in types
414 assert SFRAME_COMMIT_PACK not in types
415 assert SFRAME_OBJECT not in types
416
417
418 # ---------------------------------------------------------------------------
419 # B7 — Single commit with many objects streams correctly
420 # ---------------------------------------------------------------------------
421
422 class TestLargeSingleCommit:
423 @pytest.mark.asyncio
424 async def test_single_commit_many_objects(self, db_session: AsyncSession) -> None:
425 """A single commit referencing many objects delivers all of them."""
426 repo = await create_repo(db_session, owner="test-batch-b7")
427 repo_id = str(repo.repo_id)
428
429 n_objects = 30
430 manifest = await _make_objects(db_session, repo_id, n_objects, prefix="b7")
431 snap = await _make_snapshot(db_session, repo_id, manifest)
432 c1 = await _make_commit(
433 db_session, repo_id, snapshot_id=snap.snapshot_id, seed="b7"
434 )
435
436 from musehub.services.musehub_wire import wire_fetch_stream
437 req = WireFetchRequest(want=[c1.commit_id], have=[])
438 frames = await _collect_frames(wire_fetch_stream(db_session, repo_id, req))
439
440 received = set(_o_frame_oids(frames))
441 expected = set(manifest.values())
442 assert received == expected, (
443 f"Missing: {expected - received}, Extra: {received - expected}"
444 )
File History 1 commit
sha256:a34090cc4a394a78bd72cbbe34b08cc59525141e19135b6c0ab154f10611b9ef debug(push/stream): instrument O-frame decode path with INF… Sonnet 4.6 patch 121 days ago