"""TDD — batch streaming: wire_fetch_stream must not hold all snapshots in memory at once. The muse wire protocol mirrors git's smart HTTP pack protocol. Git streams a packfile incrementally — O(window_size) memory, not O(full_repo). We must do the same: process commits in batches, emitting O frames for new objects then a C frame for those commits, before moving to the next batch. Test plan --------- B1 O frames always precede the C frame that references them. For every commit in a C frame, every object in its snapshot manifest must have appeared in an earlier O frame. B2 Each object_id appears in at most one O frame across the whole stream. No redundant downloads. B3 A repo whose commit count exceeds _COMMIT_BATCH still delivers all commits. The stream completes correctly when multiple C frames are emitted. B4 All objects from all commits are present exactly once across all O frames. No objects missing, no duplicates. B5 Commits in different batches that share a snapshot object send that object only once (in the first batch's O frames). B6 Empty repo (no commits to fetch) produces only H + E frames — no crash. B7 Repo with a single commit spanning >_COMMIT_BATCH objects streams correctly. """ from __future__ import annotations import asyncio import zlib from datetime import datetime, timezone from typing import Any import msgpack import pytest from sqlalchemy.ext.asyncio import AsyncSession from muse.core.types import blob_id, fake_id from musehub.db import musehub_models as db from musehub.models.wire import ( SFRAME_COMMIT_PACK, SFRAME_END, SFRAME_HEADER, SFRAME_OBJECT, WireFetchRequest, ) from musehub.types.json_types import JSONObject, StrDict from tests.factories import create_repo type _OidsBySnapshot = dict[str, set[str]] # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- def _uid(seed: str) -> str: return fake_id(seed) def _now() -> datetime: return datetime.now(tz=timezone.utc) async def _make_objects( session: AsyncSession, repo_id: str, count: int, prefix: str = "obj", ) -> StrDict: """Create *count* objects and return {path: oid} manifest fragment.""" from musehub.services.musehub_wire import get_backend from sqlalchemy.dialects.postgresql import insert as pg_insert backend = get_backend() manifest: dict[str, str] = {} for i in range(count): content = f"{prefix}-{i}-content".encode() oid = blob_id(content) uri = await backend.put(oid, content) await session.execute( pg_insert(db.MusehubObject) .values( object_id=oid, path=f"{prefix}/{i}.dat", size_bytes=len(content), disk_path=uri.replace("local://", ""), storage_uri=uri, ) .on_conflict_do_nothing(index_elements=["object_id"]) ) await session.execute( pg_insert(db.MusehubObjectRef) .values(repo_id=repo_id, object_id=oid) .on_conflict_do_nothing() ) manifest[f"{prefix}/{i}.dat"] = oid await session.commit() return manifest async def _make_snapshot( session: AsyncSession, repo_id: str, manifest: dict[str, str], ) -> db.MusehubSnapshot: sid = _uid(str(sorted(manifest.items()))) snap = db.MusehubSnapshot( snapshot_id=sid, repo_id=repo_id, directories=[], manifest_blob=msgpack.packb(manifest, use_bin_type=True), entry_count=len(manifest), created_at=_now(), ) session.add(snap) await session.commit() return snap async def _make_commit( session: AsyncSession, repo_id: str, *, parent_ids: list[str] | None = None, snapshot_id: str | None = None, seed: str = "", ) -> db.MusehubCommit: commit_id = _uid(f"commit-{seed}") row = db.MusehubCommit( commit_id=commit_id, repo_id=repo_id, branch="main", parent_ids=parent_ids or [], message=f"commit {seed}", author="gabriel", timestamp=_now(), snapshot_id=snapshot_id, ) session.add(row) await session.commit() return row def _unpack_stream(chunks: list[bytes]) -> list[dict]: unpacker = msgpack.Unpacker(raw=False) for chunk in chunks: unpacker.feed(chunk) return list(unpacker) async def _collect_frames(gen: AsyncIterator[bytes]) -> list[JSONObject]: chunks: list[bytes] = [] async for chunk in gen: chunks.append(chunk) return _unpack_stream(chunks) def _o_frame_oids(frames: list[dict]) -> list[str]: """All object IDs from O frames in order.""" return [f["id"] for f in frames if f.get("t") == SFRAME_OBJECT] def _c_frame_commits(frames: list[dict]) -> list[tuple[int, dict]]: """List of (frame_index, commit_dict) for every commit in every C frame.""" result = [] for i, f in enumerate(frames): if f.get("t") == SFRAME_COMMIT_PACK: for c in f.get("commits", []): result.append((i, c)) return result def _snapshot_oids(frames: list[JSONObject]) -> _OidsBySnapshot: """Return {snapshot_id: set(oid)} from snapshot data in C frames.""" result: dict[str, set[str]] = {} for f in frames: if f.get("t") == SFRAME_COMMIT_PACK: for s in f.get("snapshots", []): sid = s.get("snapshot_id", "") oids = set(s.get("manifest", {}).values()) result[sid] = oids return result # --------------------------------------------------------------------------- # B1 — O frames precede the C frame that references them # --------------------------------------------------------------------------- class TestObjectsPrecedeCommits: @pytest.mark.asyncio async def test_objects_arrive_before_their_commit( self, db_session: AsyncSession ) -> None: """Every object in a commit's snapshot must appear in an O frame that comes before the C frame containing that commit.""" repo = await create_repo(db_session, owner="test-batch-b1") repo_id = str(repo.repo_id) manifest = await _make_objects(db_session, repo_id, 5, prefix="b1") snap = await _make_snapshot(db_session, repo_id, manifest) c1 = await _make_commit(db_session, repo_id, snapshot_id=snap.snapshot_id, seed="b1") from musehub.services.musehub_wire import wire_fetch_stream req = WireFetchRequest(want=[c1.commit_id], have=[]) frames = await _collect_frames(wire_fetch_stream(db_session, repo_id, req)) seen_oids: set[str] = set() snap_oids = _snapshot_oids(frames) for i, frame in enumerate(frames): if frame.get("t") == SFRAME_OBJECT: seen_oids.add(frame["id"]) elif frame.get("t") == SFRAME_COMMIT_PACK: for commit in frame.get("commits", []): sid = commit.get("snapshot_id", "") needed = snap_oids.get(sid, set()) missing = needed - seen_oids assert not missing, ( f"C frame at index {i} references objects not yet sent: {missing}" ) # --------------------------------------------------------------------------- # B2 — Each object appears in at most one O frame # --------------------------------------------------------------------------- class TestNoDuplicateObjects: @pytest.mark.asyncio async def test_no_object_sent_twice(self, db_session: AsyncSession) -> None: repo = await create_repo(db_session, owner="test-batch-b2") repo_id = str(repo.repo_id) manifest = await _make_objects(db_session, repo_id, 10, prefix="b2") snap = await _make_snapshot(db_session, repo_id, manifest) c1 = await _make_commit(db_session, repo_id, snapshot_id=snap.snapshot_id, seed="b2a") c2 = await _make_commit( db_session, repo_id, parent_ids=[c1.commit_id], snapshot_id=snap.snapshot_id, # same snapshot — objects must not repeat seed="b2b", ) from musehub.services.musehub_wire import wire_fetch_stream req = WireFetchRequest(want=[c2.commit_id], have=[]) frames = await _collect_frames(wire_fetch_stream(db_session, repo_id, req)) oids = _o_frame_oids(frames) assert len(oids) == len(set(oids)), ( f"Duplicate object IDs in stream: {[o for o in oids if oids.count(o) > 1]}" ) # --------------------------------------------------------------------------- # B3 — Repo with > _COMMIT_BATCH commits streams all commits across multiple C frames # --------------------------------------------------------------------------- class TestMultiBatchCompletion: @pytest.mark.asyncio async def test_all_commits_delivered_across_batches( self, db_session: AsyncSession ) -> None: """A chain of 60 commits (> _COMMIT_BATCH=50) must all appear in the stream.""" from musehub.services.musehub_wire import _COMMIT_BATCH repo = await create_repo(db_session, owner="test-batch-b3") repo_id = str(repo.repo_id) n = _COMMIT_BATCH + 10 # one more batch than the limit manifest = await _make_objects(db_session, repo_id, 3, prefix="b3") snap = await _make_snapshot(db_session, repo_id, manifest) prev_id: str | None = None commit_ids: list[str] = [] for i in range(n): c = await _make_commit( db_session, repo_id, parent_ids=[prev_id] if prev_id else [], snapshot_id=snap.snapshot_id, seed=f"b3-{i}", ) prev_id = c.commit_id commit_ids.append(c.commit_id) from musehub.services.musehub_wire import wire_fetch_stream req = WireFetchRequest(want=[prev_id], have=[]) frames = await _collect_frames(wire_fetch_stream(db_session, repo_id, req)) c_frames = [f for f in frames if f.get("t") == SFRAME_COMMIT_PACK] assert len(c_frames) >= 2, "Expected multiple C frames for a large repo" received_ids = { c["commit_id"] for cf in c_frames for c in cf.get("commits", []) } missing = set(commit_ids) - received_ids assert not missing, f"Missing commits after batch stream: {missing}" # --------------------------------------------------------------------------- # B4 — All objects from all commits present exactly once # --------------------------------------------------------------------------- class TestAllObjectsPresent: @pytest.mark.asyncio async def test_all_unique_objects_sent_exactly_once( self, db_session: AsyncSession ) -> None: repo = await create_repo(db_session, owner="test-batch-b4") repo_id = str(repo.repo_id) # Three commits, each adding distinct objects m1 = await _make_objects(db_session, repo_id, 4, prefix="b4-snap1") snap1 = await _make_snapshot(db_session, repo_id, m1) c1 = await _make_commit(db_session, repo_id, snapshot_id=snap1.snapshot_id, seed="b4-1") m2 = {**m1, **await _make_objects(db_session, repo_id, 4, prefix="b4-snap2")} snap2 = await _make_snapshot(db_session, repo_id, m2) c2 = await _make_commit( db_session, repo_id, parent_ids=[c1.commit_id], snapshot_id=snap2.snapshot_id, seed="b4-2", ) m3 = {**m2, **await _make_objects(db_session, repo_id, 4, prefix="b4-snap3")} snap3 = await _make_snapshot(db_session, repo_id, m3) c3 = await _make_commit( db_session, repo_id, parent_ids=[c2.commit_id], snapshot_id=snap3.snapshot_id, seed="b4-3", ) all_oids = set(m3.values()) # union of all three manifests from musehub.services.musehub_wire import wire_fetch_stream req = WireFetchRequest(want=[c3.commit_id], have=[]) frames = await _collect_frames(wire_fetch_stream(db_session, repo_id, req)) received = _o_frame_oids(frames) received_set = set(received) assert received_set == all_oids, ( f"Missing objects: {all_oids - received_set}\n" f"Extra objects: {received_set - all_oids}" ) assert len(received) == len(received_set), "Duplicate objects in stream" # --------------------------------------------------------------------------- # B5 — Shared objects across commit batches are sent only once # --------------------------------------------------------------------------- class TestSharedObjectsSentOnce: @pytest.mark.asyncio async def test_shared_snapshot_objects_not_duplicated( self, db_session: AsyncSession ) -> None: """Two commits in separate batches sharing the same snapshot objects must not cause those objects to be sent twice.""" from musehub.services.musehub_wire import _COMMIT_BATCH repo = await create_repo(db_session, owner="test-batch-b5") repo_id = str(repo.repo_id) shared = await _make_objects(db_session, repo_id, 5, prefix="b5-shared") snap_shared = await _make_snapshot(db_session, repo_id, shared) # Build a chain long enough to span two batches prev_id: str | None = None for i in range(_COMMIT_BATCH + 1): c = await _make_commit( db_session, repo_id, parent_ids=[prev_id] if prev_id else [], snapshot_id=snap_shared.snapshot_id, seed=f"b5-{i}", ) prev_id = c.commit_id from musehub.services.musehub_wire import wire_fetch_stream req = WireFetchRequest(want=[prev_id], have=[]) frames = await _collect_frames(wire_fetch_stream(db_session, repo_id, req)) oids = _o_frame_oids(frames) assert len(oids) == len(set(oids)), "Shared objects sent more than once" # Exactly the shared objects, no more assert set(oids) == set(shared.values()) # --------------------------------------------------------------------------- # B6 — Empty fetch (nothing to send) produces H + E only # --------------------------------------------------------------------------- class TestEmptyFetch: @pytest.mark.asyncio async def test_empty_want_produces_header_and_end( self, db_session: AsyncSession ) -> None: repo = await create_repo(db_session, owner="test-batch-b6") repo_id = str(repo.repo_id) from musehub.services.musehub_wire import wire_fetch_stream req = WireFetchRequest(want=[], have=[]) frames = await _collect_frames(wire_fetch_stream(db_session, repo_id, req)) types = [f.get("t") for f in frames] assert SFRAME_HEADER in types assert SFRAME_END in types assert SFRAME_COMMIT_PACK not in types assert SFRAME_OBJECT not in types # --------------------------------------------------------------------------- # B7 — Single commit with many objects streams correctly # --------------------------------------------------------------------------- class TestLargeSingleCommit: @pytest.mark.asyncio async def test_single_commit_many_objects(self, db_session: AsyncSession) -> None: """A single commit referencing many objects delivers all of them.""" repo = await create_repo(db_session, owner="test-batch-b7") repo_id = str(repo.repo_id) n_objects = 30 manifest = await _make_objects(db_session, repo_id, n_objects, prefix="b7") snap = await _make_snapshot(db_session, repo_id, manifest) c1 = await _make_commit( db_session, repo_id, snapshot_id=snap.snapshot_id, seed="b7" ) from musehub.services.musehub_wire import wire_fetch_stream req = WireFetchRequest(want=[c1.commit_id], have=[]) frames = await _collect_frames(wire_fetch_stream(db_session, repo_id, req)) received = set(_o_frame_oids(frames)) expected = set(manifest.values()) assert received == expected, ( f"Missing: {expected - received}, Extra: {received - expected}" )