"""Tests for MPackStreamWriter and MPackStreamReader. Coverage tiers -------------- - Writer: each write_* method produces valid msgpack with correct t tag - Writer: write_object_raw compresses and sets sz correctly - Writer: write_header includes all provenance fields - Writer: write_end echoes actual counts for integrity cross-check - Reader: feed + frames() yields complete frame dicts - Reader: handles partial chunks (incomplete msgpack across feeds) - Reader: decompress_object decodes zlib and zstd correctly - Reader: yields unknown frame types (forward compat) - Round-trip: writer → reader recovers exact payload - Protocol order: H, O×N, C, E is full valid stream - E-frame integrity: n_objects matches actual objects sent - Progress pct: float in [0.0, 100.0] - Error frame: raises note (reader yields; caller decides to raise) - Stress: 500-object stream round-trip """ from __future__ import annotations import hashlib import io import zlib import msgpack import pytest from muse.core._types import long_id def _sha(data: bytes) -> str: return long_id(hashlib.sha256(data).hexdigest()) # --------------------------------------------------------------------------- # MPackStreamWriter — write_header # --------------------------------------------------------------------------- class TestWriterHeader: def test_returns_bytes(self) -> None: from muse.core.mpack import MPackStreamWriter w = MPackStreamWriter() result = w.write_header(op="push", branch="main", n_objects=0, n_commits=0) assert isinstance(result, bytes) def test_is_valid_msgpack(self) -> None: from muse.core.mpack import MPackStreamWriter w = MPackStreamWriter() frame = msgpack.unpackb( w.write_header(op="push", branch="main", n_objects=3, n_commits=1), raw=False, ) assert isinstance(frame, dict) def test_t_field_is_H(self) -> None: from muse.core.mpack import MPackStreamWriter, FRAME_HEADER w = MPackStreamWriter() frame = msgpack.unpackb( w.write_header(op="push", branch="main", n_objects=0, n_commits=0), raw=False, ) assert frame["t"] == FRAME_HEADER def test_v_field_is_mpack_version(self) -> None: from muse.core.mpack import MPackStreamWriter, MPACK_VERSION w = MPackStreamWriter() frame = msgpack.unpackb( w.write_header(op="push", branch="main", n_objects=0, n_commits=0), raw=False, ) assert frame["v"] == MPACK_VERSION def test_op_field(self) -> None: from muse.core.mpack import MPackStreamWriter w = MPackStreamWriter() for op in ("push", "fetch", "clone"): frame = msgpack.unpackb( w.write_header(op=op, branch="main", n_objects=0, n_commits=0), raw=False, ) assert frame["op"] == op def test_provenance_agent_id(self) -> None: from muse.core.mpack import MPackStreamWriter w = MPackStreamWriter() frame = msgpack.unpackb( w.write_header( op="push", branch="main", n_objects=0, n_commits=0, agent_id="claude-code", model_id="claude-sonnet-4-6", ), raw=False, ) assert frame["agent_id"] == "claude-code" assert frame["model_id"] == "claude-sonnet-4-6" def test_domains_field(self) -> None: from muse.core.mpack import MPackStreamWriter w = MPackStreamWriter() frame = msgpack.unpackb( w.write_header( op="push", branch="main", n_objects=0, n_commits=0, domains=["code", "midi"], ), raw=False, ) assert frame["domains"] == ["code", "midi"] def test_have_list(self) -> None: from muse.core.mpack import MPackStreamWriter w = MPackStreamWriter() have = [long_id("a" * 64), long_id("b" * 64)] frame = msgpack.unpackb( w.write_header(op="push", branch="main", n_objects=0, n_commits=0, have=have), raw=False, ) assert frame["have"] == have def test_objects_bytes_advisory(self) -> None: from muse.core.mpack import MPackStreamWriter w = MPackStreamWriter() frame = msgpack.unpackb( w.write_header( op="push", branch="main", n_objects=5, n_commits=1, objects_bytes=1_000_000, ), raw=False, ) assert frame["objects_bytes"] == 1_000_000 def test_signature_bytes_present(self) -> None: from muse.core.mpack import MPackStreamWriter w = MPackStreamWriter() frame = msgpack.unpackb( w.write_header(op="push", branch="main", n_objects=0, n_commits=0), raw=False, ) assert "signature" in frame assert isinstance(frame["signature"], bytes) def test_signer_key_empty_when_no_signing(self) -> None: from muse.core.mpack import MPackStreamWriter w = MPackStreamWriter() frame = msgpack.unpackb( w.write_header(op="push", branch="main", n_objects=0, n_commits=0), raw=False, ) assert frame["signer_key"] == "" # --------------------------------------------------------------------------- # MPackStreamWriter — write_object_raw # --------------------------------------------------------------------------- class TestWriterObjectRaw: def test_returns_bytes(self) -> None: from muse.core.mpack import MPackStreamWriter w = MPackStreamWriter() raw = b"print('hello')\n" result = w.write_object_raw(object_id=_sha(raw), raw_bytes=raw) assert isinstance(result, bytes) def test_t_field_is_O(self) -> None: from muse.core.mpack import MPackStreamWriter, FRAME_OBJECT w = MPackStreamWriter() raw = b"data" frame = msgpack.unpackb( w.write_object_raw(object_id=_sha(raw), raw_bytes=raw), raw=False, ) assert frame["t"] == FRAME_OBJECT def test_id_field_is_sha256_prefixed(self) -> None: from muse.core.mpack import MPackStreamWriter w = MPackStreamWriter() raw = b"content" oid = _sha(raw) frame = msgpack.unpackb( w.write_object_raw(object_id=oid, raw_bytes=raw), raw=False, ) assert frame["id"] == oid def test_sz_equals_uncompressed_size(self) -> None: from muse.core.mpack import MPackStreamWriter w = MPackStreamWriter() raw = b"x" * 1000 frame = msgpack.unpackb( w.write_object_raw(object_id=_sha(raw), raw_bytes=raw), raw=False, ) assert frame["sz"] == 1000 def test_enc_is_zlib_or_zstd(self) -> None: from muse.core.mpack import MPackStreamWriter w = MPackStreamWriter() raw = b"content " * 100 frame = msgpack.unpackb( w.write_object_raw(object_id=_sha(raw), raw_bytes=raw), raw=False, ) assert frame["enc"] in ("zlib", "zstd") def test_force_zlib(self) -> None: from muse.core.mpack import MPackStreamWriter w = MPackStreamWriter() raw = b"data " * 50 frame = msgpack.unpackb( w.write_object_raw(object_id=_sha(raw), raw_bytes=raw, compress="zlib"), raw=False, ) assert frame["enc"] == "zlib" # Verify content is actually zlib-compressed decompressed = zlib.decompress(bytes(frame["content"])) assert decompressed == raw def test_domain_field(self) -> None: from muse.core.mpack import MPackStreamWriter w = MPackStreamWriter() raw = b"note on" frame = msgpack.unpackb( w.write_object_raw(object_id=_sha(raw), raw_bytes=raw, domain="midi"), raw=False, ) assert frame["dom"] == "midi" def test_path_field(self) -> None: from muse.core.mpack import MPackStreamWriter w = MPackStreamWriter() raw = b"content" frame = msgpack.unpackb( w.write_object_raw( object_id=_sha(raw), raw_bytes=raw, path="src/main.py", ), raw=False, ) assert frame["path"] == "src/main.py" def test_content_is_bytes(self) -> None: from muse.core.mpack import MPackStreamWriter w = MPackStreamWriter() raw = b"bytes" frame = msgpack.unpackb( w.write_object_raw(object_id=_sha(raw), raw_bytes=raw), raw=False, ) assert isinstance(frame["content"], bytes) def test_content_is_smaller_than_raw_for_compressible_data(self) -> None: from muse.core.mpack import MPackStreamWriter w = MPackStreamWriter() raw = b"aaaa" * 500 # highly compressible frame = msgpack.unpackb( w.write_object_raw(object_id=_sha(raw), raw_bytes=raw, compress="zlib"), raw=False, ) assert len(bytes(frame["content"])) < len(raw) # --------------------------------------------------------------------------- # MPackStreamWriter — write_commit_pack # --------------------------------------------------------------------------- class TestWriterCommitPack: def test_t_field_is_C(self) -> None: from muse.core.mpack import MPackStreamWriter, FRAME_COMMIT_PACK w = MPackStreamWriter() frame = msgpack.unpackb( w.write_commit_pack(commits=[], snapshots=[]), raw=False, ) assert frame["t"] == FRAME_COMMIT_PACK def test_commits_and_snapshots_preserved(self) -> None: from muse.core.mpack import MPackStreamWriter w = MPackStreamWriter() commits = [{"commit_id": long_id("a" * 64), "message": "init"}] snapshots = [{"snapshot_id": long_id("b" * 64), "manifest": {}}] frame = msgpack.unpackb( w.write_commit_pack(commits=commits, snapshots=snapshots), raw=False, ) assert len(frame["commits"]) == 1 assert len(frame["snapshots"]) == 1 # --------------------------------------------------------------------------- # MPackStreamWriter — write_end # --------------------------------------------------------------------------- class TestWriterEnd: def test_t_field_is_E(self) -> None: from muse.core.mpack import MPackStreamWriter, FRAME_END w = MPackStreamWriter() frame = msgpack.unpackb(w.write_end(n_objects=5, n_commits=2), raw=False) assert frame["t"] == FRAME_END def test_n_objects_and_n_commits_echoed(self) -> None: from muse.core.mpack import MPackStreamWriter w = MPackStreamWriter() frame = msgpack.unpackb(w.write_end(n_objects=17, n_commits=3), raw=False) assert frame["n_objects"] == 17 assert frame["n_commits"] == 3 # --------------------------------------------------------------------------- # MPackStreamWriter — response frames # --------------------------------------------------------------------------- class TestWriterResponseFrames: def test_write_progress_t_and_pct(self) -> None: from muse.core.mpack import MPackStreamWriter, FRAME_PROGRESS w = MPackStreamWriter() frame = msgpack.unpackb(w.write_progress(msg="writing objects", pct=42.5), raw=False) assert frame["t"] == FRAME_PROGRESS assert frame["pct"] == pytest.approx(42.5) assert frame["msg"] == "writing objects" def test_write_error_t_msg_code(self) -> None: from muse.core.mpack import MPackStreamWriter, FRAME_ERROR w = MPackStreamWriter() frame = msgpack.unpackb(w.write_error(msg="not allowed", code=403), raw=False) assert frame["t"] == FRAME_ERROR assert frame["msg"] == "not allowed" assert frame["code"] == 403 def test_write_result_ok_true(self) -> None: from muse.core.mpack import MPackStreamWriter, FRAME_RESULT w = MPackStreamWriter() heads = {"main": long_id("a" * 64)} frame = msgpack.unpackb( w.write_result(ok=True, msg="pushed", heads=heads, head=long_id("a" * 64)), raw=False, ) assert frame["t"] == FRAME_RESULT assert frame["ok"] is True assert frame["heads"] == heads def test_write_ack_ready(self) -> None: from muse.core.mpack import MPackStreamWriter, FRAME_ACK w = MPackStreamWriter() frame = msgpack.unpackb( w.write_ack(ack=[long_id("a" * 64)], common=long_id("a" * 64), ready=True), raw=False, ) assert frame["t"] == FRAME_ACK assert frame["ready"] is True # --------------------------------------------------------------------------- # MPackStreamReader — feed + frames # --------------------------------------------------------------------------- class TestReaderBasic: def test_returns_frame_dicts(self) -> None: from muse.core.mpack import MPackStreamWriter, MPackStreamReader w = MPackStreamWriter() r = MPackStreamReader() r.feed(w.write_header(op="push", branch="main", n_objects=0, n_commits=0)) frames = list(r.frames()) assert len(frames) == 1 assert frames[0]["t"] == "H" def test_handles_partial_chunks(self) -> None: from muse.core.mpack import MPackStreamWriter, MPackStreamReader w = MPackStreamWriter() raw = b"partial-feed data" encoded = w.write_object_raw(object_id=_sha(raw), raw_bytes=raw) r = MPackStreamReader() # Feed one byte at a time frames = [] for byte in encoded: r.feed(bytes([byte])) frames.extend(r.frames()) assert len(frames) == 1 assert frames[0]["t"] == "O" def test_multiple_frames_in_one_feed(self) -> None: from muse.core.mpack import MPackStreamWriter, MPackStreamReader w = MPackStreamWriter() body = ( w.write_header(op="push", branch="main", n_objects=1, n_commits=1) + w.write_end(n_objects=1, n_commits=1) ) r = MPackStreamReader() r.feed(body) frames = list(r.frames()) assert len(frames) == 2 assert frames[0]["t"] == "H" assert frames[1]["t"] == "E" def test_unknown_frame_types_are_yielded(self) -> None: """Unknown t values must be yielded, not swallowed — forward compat.""" from muse.core.mpack import MPackStreamReader unknown = msgpack.packb({"t": "Z", "data": "future feature"}, use_bin_type=True) r = MPackStreamReader() r.feed(unknown) frames = list(r.frames()) assert len(frames) == 1 assert frames[0]["t"] == "Z" def test_non_dict_frames_are_skipped(self) -> None: from muse.core.mpack import MPackStreamReader garbage = msgpack.packb([1, 2, 3], use_bin_type=True) r = MPackStreamReader() r.feed(garbage) frames = list(r.frames()) assert frames == [] # --------------------------------------------------------------------------- # MPackStreamReader — decompress_object # --------------------------------------------------------------------------- class TestReaderDecompressObject: def test_decompress_zlib(self) -> None: from muse.core.mpack import MPackStreamReader raw = b"hello world " * 100 compressed = zlib.compress(raw) frame = {"t": "O", "content": compressed, "enc": "zlib", "sz": len(raw)} r = MPackStreamReader() result = r.decompress_object(frame) assert result == raw def test_decompress_raw(self) -> None: from muse.core.mpack import MPackStreamReader raw = b"uncompressed" frame = {"t": "O", "content": raw, "enc": "raw", "sz": len(raw)} r = MPackStreamReader() result = r.decompress_object(frame) assert result == raw def test_delta_raises_without_base(self) -> None: from muse.core.mpack import MPackStreamReader frame = {"t": "O", "content": b"delta", "enc": "delta+zlib", "sz": 100} r = MPackStreamReader() with pytest.raises(ValueError, match="delta"): r.decompress_object(frame) # --------------------------------------------------------------------------- # Round-trip: writer → reader # --------------------------------------------------------------------------- class TestRoundTrip: def test_header_round_trip(self) -> None: from muse.core.mpack import MPackStreamWriter, MPackStreamReader w = MPackStreamWriter() r = MPackStreamReader() r.feed(w.write_header( op="push", branch="dev", n_objects=3, n_commits=1, agent_id="claude-code", model_id="claude-sonnet-4-6", domains=["code"], objects_bytes=9000, )) frames = list(r.frames()) h = frames[0] assert h["op"] == "push" assert h["branch"] == "dev" assert h["n_objects"] == 3 assert h["agent_id"] == "claude-code" assert h["domains"] == ["code"] assert h["objects_bytes"] == 9000 def test_object_round_trip_zlib(self) -> None: from muse.core.mpack import MPackStreamWriter, MPackStreamReader w = MPackStreamWriter() r = MPackStreamReader() raw = b"def hello(): pass\n" * 20 r.feed(w.write_object_raw( object_id=_sha(raw), raw_bytes=raw, path="hello.py", domain="code", compress="zlib", )) frames = list(r.frames()) assert frames[0]["sz"] == len(raw) decompressed = r.decompress_object(frames[0]) assert decompressed == raw def test_commit_pack_round_trip(self) -> None: from muse.core.mpack import MPackStreamWriter, MPackStreamReader w = MPackStreamWriter() r = MPackStreamReader() commits = [{"commit_id": long_id("c" * 64), "message": "feat: add MPack"}] snaps = [{"snapshot_id": long_id("d" * 64), "manifest": {"a.py": long_id("e" * 64)}}] r.feed(w.write_commit_pack(commits=commits, snapshots=snaps)) frames = list(r.frames()) assert frames[0]["commits"][0]["message"] == "feat: add MPack" assert frames[0]["snapshots"][0]["manifest"]["a.py"] == long_id("e" * 64) def test_full_push_stream_round_trip(self) -> None: from muse.core.mpack import MPackStreamWriter, MPackStreamReader, FRAME_END w = MPackStreamWriter() r = MPackStreamReader() objects = [f"object {i}\n".encode() * 10 for i in range(5)] body = w.write_header(op="push", branch="main", n_objects=5, n_commits=1) for raw in objects: body += w.write_object_raw( object_id=_sha(raw), raw_bytes=raw, compress="zlib", ) body += w.write_commit_pack(commits=[], snapshots=[]) body += w.write_end(n_objects=5, n_commits=1) r.feed(body) frames = list(r.frames()) tags = [f["t"] for f in frames] assert tags[0] == "H" assert tags[1:-2] == ["O"] * 5 assert tags[-2] == "C" assert tags[-1] == FRAME_END # E frame integrity cross-check end = frames[-1] assert end["n_objects"] == 5 assert end["n_commits"] == 1 def test_end_frame_n_objects_matches_actual(self) -> None: from muse.core.mpack import MPackStreamWriter, MPackStreamReader w = MPackStreamWriter() r = MPackStreamReader() n = 7 body = w.write_header(op="push", branch="main", n_objects=n, n_commits=0) for i in range(n): raw = f"data{i}".encode() body += w.write_object_raw(object_id=_sha(raw), raw_bytes=raw) body += w.write_commit_pack(commits=[], snapshots=[]) body += w.write_end(n_objects=n, n_commits=0) r.feed(body) frames = list(r.frames()) end = frames[-1] assert end["n_objects"] == n # --------------------------------------------------------------------------- # Stress: 500-object stream # --------------------------------------------------------------------------- class TestStress: def test_500_object_stream(self) -> None: from muse.core.mpack import MPackStreamWriter, MPackStreamReader w = MPackStreamWriter() r = MPackStreamReader() n = 500 body = w.write_header(op="push", branch="main", n_objects=n, n_commits=1) for i in range(n): raw = f"content of object {i}\n".encode() * 5 body += w.write_object_raw(object_id=_sha(raw), raw_bytes=raw, compress="zlib") body += w.write_commit_pack(commits=[], snapshots=[]) body += w.write_end(n_objects=n, n_commits=1) r.feed(body) frames = list(r.frames()) object_frames = [f for f in frames if f["t"] == "O"] assert len(object_frames) == n # Every object frame must have sz and enc fields for f in object_frames: assert "sz" in f and f["sz"] > 0 assert "enc" in f and f["enc"] in ("zlib", "zstd")