gabriel / muse public
test_mpack_stream.py python
565 lines 20.8 KB
Raw
sha256:88ac91129873e6a496e9189515aa690eb893ae25d69c8f72af141a2be5068eb3 docs: docstring sprint contract→find-symbol — idiomatic run… Sonnet 4.6 patch 142 days ago
1 """Tests for MPackStreamWriter and MPackStreamReader.
2
3 Coverage tiers
4 --------------
5 - Writer: each write_* method produces valid msgpack with correct t tag
6 - Writer: write_object_raw compresses and sets sz correctly
7 - Writer: write_header includes all provenance fields
8 - Writer: write_end echoes actual counts for integrity cross-check
9 - Reader: feed + frames() yields complete frame dicts
10 - Reader: handles partial chunks (incomplete msgpack across feeds)
11 - Reader: decompress_object decodes zlib and zstd correctly
12 - Reader: yields unknown frame types (forward compat)
13 - Round-trip: writer → reader recovers exact payload
14 - Protocol order: H, O×N, C, E is full valid stream
15 - E-frame integrity: n_objects matches actual objects sent
16 - Progress pct: float in [0.0, 100.0]
17 - Error frame: raises note (reader yields; caller decides to raise)
18 - Stress: 500-object stream round-trip
19 """
20 from __future__ import annotations
21
22 import hashlib
23 import io
24 import zlib
25
26 import msgpack
27 import pytest
28
29 from muse.core._types import long_id
30
31
32 def _sha(data: bytes) -> str:
33 return long_id(hashlib.sha256(data).hexdigest())
34
35
36 # ---------------------------------------------------------------------------
37 # MPackStreamWriter — write_header
38 # ---------------------------------------------------------------------------
39
40
41 class TestWriterHeader:
42 def test_returns_bytes(self) -> None:
43 from muse.core.mpack import MPackStreamWriter
44 w = MPackStreamWriter()
45 result = w.write_header(op="push", branch="main", n_objects=0, n_commits=0)
46 assert isinstance(result, bytes)
47
48 def test_is_valid_msgpack(self) -> None:
49 from muse.core.mpack import MPackStreamWriter
50 w = MPackStreamWriter()
51 frame = msgpack.unpackb(
52 w.write_header(op="push", branch="main", n_objects=3, n_commits=1),
53 raw=False,
54 )
55 assert isinstance(frame, dict)
56
57 def test_t_field_is_H(self) -> None:
58 from muse.core.mpack import MPackStreamWriter, FRAME_HEADER
59 w = MPackStreamWriter()
60 frame = msgpack.unpackb(
61 w.write_header(op="push", branch="main", n_objects=0, n_commits=0),
62 raw=False,
63 )
64 assert frame["t"] == FRAME_HEADER
65
66 def test_v_field_is_mpack_version(self) -> None:
67 from muse.core.mpack import MPackStreamWriter, MPACK_VERSION
68 w = MPackStreamWriter()
69 frame = msgpack.unpackb(
70 w.write_header(op="push", branch="main", n_objects=0, n_commits=0),
71 raw=False,
72 )
73 assert frame["v"] == MPACK_VERSION
74
75 def test_op_field(self) -> None:
76 from muse.core.mpack import MPackStreamWriter
77 w = MPackStreamWriter()
78 for op in ("push", "fetch", "clone"):
79 frame = msgpack.unpackb(
80 w.write_header(op=op, branch="main", n_objects=0, n_commits=0),
81 raw=False,
82 )
83 assert frame["op"] == op
84
85 def test_provenance_agent_id(self) -> None:
86 from muse.core.mpack import MPackStreamWriter
87 w = MPackStreamWriter()
88 frame = msgpack.unpackb(
89 w.write_header(
90 op="push", branch="main", n_objects=0, n_commits=0,
91 agent_id="claude-code", model_id="claude-sonnet-4-6",
92 ),
93 raw=False,
94 )
95 assert frame["agent_id"] == "claude-code"
96 assert frame["model_id"] == "claude-sonnet-4-6"
97
98 def test_domains_field(self) -> None:
99 from muse.core.mpack import MPackStreamWriter
100 w = MPackStreamWriter()
101 frame = msgpack.unpackb(
102 w.write_header(
103 op="push", branch="main", n_objects=0, n_commits=0,
104 domains=["code", "midi"],
105 ),
106 raw=False,
107 )
108 assert frame["domains"] == ["code", "midi"]
109
110 def test_have_list(self) -> None:
111 from muse.core.mpack import MPackStreamWriter
112 w = MPackStreamWriter()
113 have = [long_id("a" * 64), long_id("b" * 64)]
114 frame = msgpack.unpackb(
115 w.write_header(op="push", branch="main", n_objects=0, n_commits=0, have=have),
116 raw=False,
117 )
118 assert frame["have"] == have
119
120 def test_objects_bytes_advisory(self) -> None:
121 from muse.core.mpack import MPackStreamWriter
122 w = MPackStreamWriter()
123 frame = msgpack.unpackb(
124 w.write_header(
125 op="push", branch="main", n_objects=5, n_commits=1, objects_bytes=1_000_000,
126 ),
127 raw=False,
128 )
129 assert frame["objects_bytes"] == 1_000_000
130
131 def test_signature_bytes_present(self) -> None:
132 from muse.core.mpack import MPackStreamWriter
133 w = MPackStreamWriter()
134 frame = msgpack.unpackb(
135 w.write_header(op="push", branch="main", n_objects=0, n_commits=0),
136 raw=False,
137 )
138 assert "signature" in frame
139 assert isinstance(frame["signature"], bytes)
140
141 def test_signer_key_empty_when_no_signing(self) -> None:
142 from muse.core.mpack import MPackStreamWriter
143 w = MPackStreamWriter()
144 frame = msgpack.unpackb(
145 w.write_header(op="push", branch="main", n_objects=0, n_commits=0),
146 raw=False,
147 )
148 assert frame["signer_key"] == ""
149
150
151 # ---------------------------------------------------------------------------
152 # MPackStreamWriter — write_object_raw
153 # ---------------------------------------------------------------------------
154
155
156 class TestWriterObjectRaw:
157 def test_returns_bytes(self) -> None:
158 from muse.core.mpack import MPackStreamWriter
159 w = MPackStreamWriter()
160 raw = b"print('hello')\n"
161 result = w.write_object_raw(object_id=_sha(raw), raw_bytes=raw)
162 assert isinstance(result, bytes)
163
164 def test_t_field_is_O(self) -> None:
165 from muse.core.mpack import MPackStreamWriter, FRAME_OBJECT
166 w = MPackStreamWriter()
167 raw = b"data"
168 frame = msgpack.unpackb(
169 w.write_object_raw(object_id=_sha(raw), raw_bytes=raw),
170 raw=False,
171 )
172 assert frame["t"] == FRAME_OBJECT
173
174 def test_id_field_is_sha256_prefixed(self) -> None:
175 from muse.core.mpack import MPackStreamWriter
176 w = MPackStreamWriter()
177 raw = b"content"
178 oid = _sha(raw)
179 frame = msgpack.unpackb(
180 w.write_object_raw(object_id=oid, raw_bytes=raw),
181 raw=False,
182 )
183 assert frame["id"] == oid
184
185 def test_sz_equals_uncompressed_size(self) -> None:
186 from muse.core.mpack import MPackStreamWriter
187 w = MPackStreamWriter()
188 raw = b"x" * 1000
189 frame = msgpack.unpackb(
190 w.write_object_raw(object_id=_sha(raw), raw_bytes=raw),
191 raw=False,
192 )
193 assert frame["sz"] == 1000
194
195 def test_enc_is_zlib_or_zstd(self) -> None:
196 from muse.core.mpack import MPackStreamWriter
197 w = MPackStreamWriter()
198 raw = b"content " * 100
199 frame = msgpack.unpackb(
200 w.write_object_raw(object_id=_sha(raw), raw_bytes=raw),
201 raw=False,
202 )
203 assert frame["enc"] in ("zlib", "zstd")
204
205 def test_force_zlib(self) -> None:
206 from muse.core.mpack import MPackStreamWriter
207 w = MPackStreamWriter()
208 raw = b"data " * 50
209 frame = msgpack.unpackb(
210 w.write_object_raw(object_id=_sha(raw), raw_bytes=raw, compress="zlib"),
211 raw=False,
212 )
213 assert frame["enc"] == "zlib"
214 # Verify content is actually zlib-compressed
215 decompressed = zlib.decompress(bytes(frame["content"]))
216 assert decompressed == raw
217
218 def test_domain_field(self) -> None:
219 from muse.core.mpack import MPackStreamWriter
220 w = MPackStreamWriter()
221 raw = b"note on"
222 frame = msgpack.unpackb(
223 w.write_object_raw(object_id=_sha(raw), raw_bytes=raw, domain="midi"),
224 raw=False,
225 )
226 assert frame["dom"] == "midi"
227
228 def test_path_field(self) -> None:
229 from muse.core.mpack import MPackStreamWriter
230 w = MPackStreamWriter()
231 raw = b"content"
232 frame = msgpack.unpackb(
233 w.write_object_raw(
234 object_id=_sha(raw), raw_bytes=raw, path="src/main.py",
235 ),
236 raw=False,
237 )
238 assert frame["path"] == "src/main.py"
239
240 def test_content_is_bytes(self) -> None:
241 from muse.core.mpack import MPackStreamWriter
242 w = MPackStreamWriter()
243 raw = b"bytes"
244 frame = msgpack.unpackb(
245 w.write_object_raw(object_id=_sha(raw), raw_bytes=raw),
246 raw=False,
247 )
248 assert isinstance(frame["content"], bytes)
249
250 def test_content_is_smaller_than_raw_for_compressible_data(self) -> None:
251 from muse.core.mpack import MPackStreamWriter
252 w = MPackStreamWriter()
253 raw = b"aaaa" * 500 # highly compressible
254 frame = msgpack.unpackb(
255 w.write_object_raw(object_id=_sha(raw), raw_bytes=raw, compress="zlib"),
256 raw=False,
257 )
258 assert len(bytes(frame["content"])) < len(raw)
259
260
261 # ---------------------------------------------------------------------------
262 # MPackStreamWriter — write_commit_pack
263 # ---------------------------------------------------------------------------
264
265
266 class TestWriterCommitPack:
267 def test_t_field_is_C(self) -> None:
268 from muse.core.mpack import MPackStreamWriter, FRAME_COMMIT_PACK
269 w = MPackStreamWriter()
270 frame = msgpack.unpackb(
271 w.write_commit_pack(commits=[], snapshots=[]),
272 raw=False,
273 )
274 assert frame["t"] == FRAME_COMMIT_PACK
275
276 def test_commits_and_snapshots_preserved(self) -> None:
277 from muse.core.mpack import MPackStreamWriter
278 w = MPackStreamWriter()
279 commits = [{"commit_id": long_id("a" * 64), "message": "init"}]
280 snapshots = [{"snapshot_id": long_id("b" * 64), "manifest": {}}]
281 frame = msgpack.unpackb(
282 w.write_commit_pack(commits=commits, snapshots=snapshots),
283 raw=False,
284 )
285 assert len(frame["commits"]) == 1
286 assert len(frame["snapshots"]) == 1
287
288
289 # ---------------------------------------------------------------------------
290 # MPackStreamWriter — write_end
291 # ---------------------------------------------------------------------------
292
293
294 class TestWriterEnd:
295 def test_t_field_is_E(self) -> None:
296 from muse.core.mpack import MPackStreamWriter, FRAME_END
297 w = MPackStreamWriter()
298 frame = msgpack.unpackb(w.write_end(n_objects=5, n_commits=2), raw=False)
299 assert frame["t"] == FRAME_END
300
301 def test_n_objects_and_n_commits_echoed(self) -> None:
302 from muse.core.mpack import MPackStreamWriter
303 w = MPackStreamWriter()
304 frame = msgpack.unpackb(w.write_end(n_objects=17, n_commits=3), raw=False)
305 assert frame["n_objects"] == 17
306 assert frame["n_commits"] == 3
307
308
309 # ---------------------------------------------------------------------------
310 # MPackStreamWriter — response frames
311 # ---------------------------------------------------------------------------
312
313
314 class TestWriterResponseFrames:
315 def test_write_progress_t_and_pct(self) -> None:
316 from muse.core.mpack import MPackStreamWriter, FRAME_PROGRESS
317 w = MPackStreamWriter()
318 frame = msgpack.unpackb(w.write_progress(msg="writing objects", pct=42.5), raw=False)
319 assert frame["t"] == FRAME_PROGRESS
320 assert frame["pct"] == pytest.approx(42.5)
321 assert frame["msg"] == "writing objects"
322
323 def test_write_error_t_msg_code(self) -> None:
324 from muse.core.mpack import MPackStreamWriter, FRAME_ERROR
325 w = MPackStreamWriter()
326 frame = msgpack.unpackb(w.write_error(msg="not allowed", code=403), raw=False)
327 assert frame["t"] == FRAME_ERROR
328 assert frame["msg"] == "not allowed"
329 assert frame["code"] == 403
330
331 def test_write_result_ok_true(self) -> None:
332 from muse.core.mpack import MPackStreamWriter, FRAME_RESULT
333 w = MPackStreamWriter()
334 heads = {"main": long_id("a" * 64)}
335 frame = msgpack.unpackb(
336 w.write_result(ok=True, msg="pushed", heads=heads, head=long_id("a" * 64)),
337 raw=False,
338 )
339 assert frame["t"] == FRAME_RESULT
340 assert frame["ok"] is True
341 assert frame["heads"] == heads
342
343 def test_write_ack_ready(self) -> None:
344 from muse.core.mpack import MPackStreamWriter, FRAME_ACK
345 w = MPackStreamWriter()
346 frame = msgpack.unpackb(
347 w.write_ack(ack=[long_id("a" * 64)], common=long_id("a" * 64), ready=True),
348 raw=False,
349 )
350 assert frame["t"] == FRAME_ACK
351 assert frame["ready"] is True
352
353
354 # ---------------------------------------------------------------------------
355 # MPackStreamReader — feed + frames
356 # ---------------------------------------------------------------------------
357
358
359 class TestReaderBasic:
360 def test_returns_frame_dicts(self) -> None:
361 from muse.core.mpack import MPackStreamWriter, MPackStreamReader
362 w = MPackStreamWriter()
363 r = MPackStreamReader()
364 r.feed(w.write_header(op="push", branch="main", n_objects=0, n_commits=0))
365 frames = list(r.frames())
366 assert len(frames) == 1
367 assert frames[0]["t"] == "H"
368
369 def test_handles_partial_chunks(self) -> None:
370 from muse.core.mpack import MPackStreamWriter, MPackStreamReader
371 w = MPackStreamWriter()
372 raw = b"partial-feed data"
373 encoded = w.write_object_raw(object_id=_sha(raw), raw_bytes=raw)
374 r = MPackStreamReader()
375 # Feed one byte at a time
376 frames = []
377 for byte in encoded:
378 r.feed(bytes([byte]))
379 frames.extend(r.frames())
380 assert len(frames) == 1
381 assert frames[0]["t"] == "O"
382
383 def test_multiple_frames_in_one_feed(self) -> None:
384 from muse.core.mpack import MPackStreamWriter, MPackStreamReader
385 w = MPackStreamWriter()
386 body = (
387 w.write_header(op="push", branch="main", n_objects=1, n_commits=1)
388 + w.write_end(n_objects=1, n_commits=1)
389 )
390 r = MPackStreamReader()
391 r.feed(body)
392 frames = list(r.frames())
393 assert len(frames) == 2
394 assert frames[0]["t"] == "H"
395 assert frames[1]["t"] == "E"
396
397 def test_unknown_frame_types_are_yielded(self) -> None:
398 """Unknown t values must be yielded, not swallowed — forward compat."""
399 from muse.core.mpack import MPackStreamReader
400 unknown = msgpack.packb({"t": "Z", "data": "future feature"}, use_bin_type=True)
401 r = MPackStreamReader()
402 r.feed(unknown)
403 frames = list(r.frames())
404 assert len(frames) == 1
405 assert frames[0]["t"] == "Z"
406
407 def test_non_dict_frames_are_skipped(self) -> None:
408 from muse.core.mpack import MPackStreamReader
409 garbage = msgpack.packb([1, 2, 3], use_bin_type=True)
410 r = MPackStreamReader()
411 r.feed(garbage)
412 frames = list(r.frames())
413 assert frames == []
414
415
416 # ---------------------------------------------------------------------------
417 # MPackStreamReader — decompress_object
418 # ---------------------------------------------------------------------------
419
420
421 class TestReaderDecompressObject:
422 def test_decompress_zlib(self) -> None:
423 from muse.core.mpack import MPackStreamReader
424 raw = b"hello world " * 100
425 compressed = zlib.compress(raw)
426 frame = {"t": "O", "content": compressed, "enc": "zlib", "sz": len(raw)}
427 r = MPackStreamReader()
428 result = r.decompress_object(frame)
429 assert result == raw
430
431 def test_decompress_raw(self) -> None:
432 from muse.core.mpack import MPackStreamReader
433 raw = b"uncompressed"
434 frame = {"t": "O", "content": raw, "enc": "raw", "sz": len(raw)}
435 r = MPackStreamReader()
436 result = r.decompress_object(frame)
437 assert result == raw
438
439 def test_delta_raises_without_base(self) -> None:
440 from muse.core.mpack import MPackStreamReader
441 frame = {"t": "O", "content": b"delta", "enc": "delta+zlib", "sz": 100}
442 r = MPackStreamReader()
443 with pytest.raises(ValueError, match="delta"):
444 r.decompress_object(frame)
445
446
447 # ---------------------------------------------------------------------------
448 # Round-trip: writer → reader
449 # ---------------------------------------------------------------------------
450
451
452 class TestRoundTrip:
453 def test_header_round_trip(self) -> None:
454 from muse.core.mpack import MPackStreamWriter, MPackStreamReader
455 w = MPackStreamWriter()
456 r = MPackStreamReader()
457 r.feed(w.write_header(
458 op="push", branch="dev", n_objects=3, n_commits=1,
459 agent_id="claude-code", model_id="claude-sonnet-4-6",
460 domains=["code"], objects_bytes=9000,
461 ))
462 frames = list(r.frames())
463 h = frames[0]
464 assert h["op"] == "push"
465 assert h["branch"] == "dev"
466 assert h["n_objects"] == 3
467 assert h["agent_id"] == "claude-code"
468 assert h["domains"] == ["code"]
469 assert h["objects_bytes"] == 9000
470
471 def test_object_round_trip_zlib(self) -> None:
472 from muse.core.mpack import MPackStreamWriter, MPackStreamReader
473 w = MPackStreamWriter()
474 r = MPackStreamReader()
475 raw = b"def hello(): pass\n" * 20
476 r.feed(w.write_object_raw(
477 object_id=_sha(raw), raw_bytes=raw, path="hello.py", domain="code",
478 compress="zlib",
479 ))
480 frames = list(r.frames())
481 assert frames[0]["sz"] == len(raw)
482 decompressed = r.decompress_object(frames[0])
483 assert decompressed == raw
484
485 def test_commit_pack_round_trip(self) -> None:
486 from muse.core.mpack import MPackStreamWriter, MPackStreamReader
487 w = MPackStreamWriter()
488 r = MPackStreamReader()
489 commits = [{"commit_id": long_id("c" * 64), "message": "feat: add MPack"}]
490 snaps = [{"snapshot_id": long_id("d" * 64), "manifest": {"a.py": long_id("e" * 64)}}]
491 r.feed(w.write_commit_pack(commits=commits, snapshots=snaps))
492 frames = list(r.frames())
493 assert frames[0]["commits"][0]["message"] == "feat: add MPack"
494 assert frames[0]["snapshots"][0]["manifest"]["a.py"] == long_id("e" * 64)
495
496 def test_full_push_stream_round_trip(self) -> None:
497 from muse.core.mpack import MPackStreamWriter, MPackStreamReader, FRAME_END
498 w = MPackStreamWriter()
499 r = MPackStreamReader()
500
501 objects = [f"object {i}\n".encode() * 10 for i in range(5)]
502 body = w.write_header(op="push", branch="main", n_objects=5, n_commits=1)
503 for raw in objects:
504 body += w.write_object_raw(
505 object_id=_sha(raw), raw_bytes=raw, compress="zlib",
506 )
507 body += w.write_commit_pack(commits=[], snapshots=[])
508 body += w.write_end(n_objects=5, n_commits=1)
509
510 r.feed(body)
511 frames = list(r.frames())
512 tags = [f["t"] for f in frames]
513 assert tags[0] == "H"
514 assert tags[1:-2] == ["O"] * 5
515 assert tags[-2] == "C"
516 assert tags[-1] == FRAME_END
517
518 # E frame integrity cross-check
519 end = frames[-1]
520 assert end["n_objects"] == 5
521 assert end["n_commits"] == 1
522
523 def test_end_frame_n_objects_matches_actual(self) -> None:
524 from muse.core.mpack import MPackStreamWriter, MPackStreamReader
525 w = MPackStreamWriter()
526 r = MPackStreamReader()
527 n = 7
528 body = w.write_header(op="push", branch="main", n_objects=n, n_commits=0)
529 for i in range(n):
530 raw = f"data{i}".encode()
531 body += w.write_object_raw(object_id=_sha(raw), raw_bytes=raw)
532 body += w.write_commit_pack(commits=[], snapshots=[])
533 body += w.write_end(n_objects=n, n_commits=0)
534 r.feed(body)
535 frames = list(r.frames())
536 end = frames[-1]
537 assert end["n_objects"] == n
538
539
540 # ---------------------------------------------------------------------------
541 # Stress: 500-object stream
542 # ---------------------------------------------------------------------------
543
544
545 class TestStress:
546 def test_500_object_stream(self) -> None:
547 from muse.core.mpack import MPackStreamWriter, MPackStreamReader
548 w = MPackStreamWriter()
549 r = MPackStreamReader()
550 n = 500
551 body = w.write_header(op="push", branch="main", n_objects=n, n_commits=1)
552 for i in range(n):
553 raw = f"content of object {i}\n".encode() * 5
554 body += w.write_object_raw(object_id=_sha(raw), raw_bytes=raw, compress="zlib")
555 body += w.write_commit_pack(commits=[], snapshots=[])
556 body += w.write_end(n_objects=n, n_commits=1)
557
558 r.feed(body)
559 frames = list(r.frames())
560 object_frames = [f for f in frames if f["t"] == "O"]
561 assert len(object_frames) == n
562 # Every object frame must have sz and enc fields
563 for f in object_frames:
564 assert "sz" in f and f["sz"] > 0
565 assert "enc" in f and f["enc"] in ("zlib", "zstd")
File History 1 commit
sha256:88ac91129873e6a496e9189515aa690eb893ae25d69c8f72af141a2be5068eb3 docs: docstring sprint contract→find-symbol — idiomatic run… Sonnet 4.6 patch 142 days ago