gabriel / muse public
test_snapshot_schema_version_and_compression.py python
480 lines 20.7 KB
Raw
sha256:51ce277f663e01a43eaffbe77509b1de7ac2d4251b55d23306304bcdeb92c90d feat(pack): delta-encode snapshots in MPackBundle wire format Sonnet 4.6 minor ⚠ breaking 121 days ago
1 """Tests for snapshot schema_version field and zstd at-rest compression.
2
3 Every new snapshot file is written as ``zstd(msgpack(data))`` when the packed
4 payload exceeds ``_ZSTD_COMPRESS_THRESHOLD`` bytes. Smaller snapshots stay as
5 raw msgpack — no overhead for tiny repos. Detection is self-describing via the
6 4-byte zstd magic ``\\x28\\xb5\\x2f\\xfd`` at the start of the file, so old
7 uncompressed files remain fully readable without any migration.
8
9 ``schema_version`` (integer, currently 1) is stored in each snapshot record as
10 metadata. It is intentionally excluded from the snapshot-ID hash — the hash
11 captures only content (manifest paths + object IDs + directories). This lets
12 the schema version evolve (e.g. when the Rust port lands) without invalidating
13 any existing snapshot ID.
14
15 Seven-tier coverage
16 -------------------
17 - Unit — constants, zstd helpers, schema_version field contract
18 - Integration — write/read roundtrip with schema_version and compression
19 - E2E — full CLI: ``muse snapshot create`` stores compressed file on disk
20 - Stress — 1 000-file manifest compresses and decompresses without error
21 - State — pre-compression (uncompressed) snapshots are still readable
22 - Integrity — ``_verify_snapshot_id`` passes on compressed snapshots;
23 schema_version cannot alter the content hash
24 - Performance — 1 000-file roundtrip completes within 2 s
25 - Security — zstd "bomb" that expands beyond MAX_MSGPACK_BYTES is rejected
26 """
27
28 from __future__ import annotations
29
30 import datetime
31 import pathlib
32 import time
33
34 import msgpack
35 import pytest
36
37 from muse.core.snapshot import compute_snapshot_id
38 from muse.core.store import (
39 MAX_MSGPACK_BYTES,
40 SnapshotRecord,
41 read_snapshot,
42 snapshot_path,
43 write_snapshot,
44 )
45 from muse.core.types import content_hash, long_id
46 from muse.core.paths import muse_dir, snapshots_dir
47
48
49 # ---------------------------------------------------------------------------
50 # Helpers shared across tiers
51 # ---------------------------------------------------------------------------
52
53
54 def _obj_id(n: int) -> str:
55 return long_id(f"{n:064x}")
56
57
58 def _make_snapshot(n_files: int = 5, note: str = "") -> SnapshotRecord:
59 manifest = {f"src/file_{i:04d}.py": _obj_id(i) for i in range(n_files)}
60 snap_id = compute_snapshot_id(manifest)
61 return SnapshotRecord(
62 snapshot_id=snap_id,
63 manifest=manifest,
64 note=note,
65 )
66
67
68 def _init_repo(tmp_path: pathlib.Path) -> pathlib.Path:
69 """Minimal .muse/ tree — just enough for snapshot read/write."""
70 muse = muse_dir(tmp_path)
71 (muse / "snapshots").mkdir(parents=True)
72 return tmp_path
73
74
75 # ---------------------------------------------------------------------------
76 # Tier 1 — Unit: constants and zstd helper contract
77 # ---------------------------------------------------------------------------
78
79
80 class TestConstants:
81 def test_snapshot_schema_version_is_int(self) -> None:
82 """_SNAPSHOT_SCHEMA_VERSION must be a plain int, not str or float."""
83 from muse.core.store import _SNAPSHOT_SCHEMA_VERSION
84 assert isinstance(_SNAPSHOT_SCHEMA_VERSION, int)
85
86 def test_snapshot_schema_version_is_one(self) -> None:
87 """Current schema version is 1 — bump only on breaking layout changes."""
88 from muse.core.store import _SNAPSHOT_SCHEMA_VERSION
89 assert _SNAPSHOT_SCHEMA_VERSION == 1
90
91 def test_zstd_magic_is_correct(self) -> None:
92 """The 4-byte zstd frame magic must match the zstd specification."""
93 from muse.core.store import _ZSTD_MAGIC
94 assert _ZSTD_MAGIC == b"\x28\xb5\x2f\xfd"
95
96 def test_compress_threshold_is_positive(self) -> None:
97 from muse.core.store import _ZSTD_COMPRESS_THRESHOLD
98 assert _ZSTD_COMPRESS_THRESHOLD > 0
99
100 def test_compress_threshold_is_reasonable(self) -> None:
101 """Threshold must be large enough that single-file snapshots are not compressed."""
102 from muse.core.store import _ZSTD_COMPRESS_THRESHOLD
103 assert _ZSTD_COMPRESS_THRESHOLD >= 1024
104
105
106 class TestZstdHelpers:
107 def test_zstd_roundtrip(self) -> None:
108 """compress → decompress_if_needed must return the original bytes exactly."""
109 from muse.core.store import _zstd_compress, zstd_decompress_if_needed
110 original = b"hello " * 1_000
111 compressed = _zstd_compress(original)
112 recovered = zstd_decompress_if_needed(compressed)
113 assert recovered == original
114
115 def test_compressed_output_starts_with_magic(self) -> None:
116 """zstd output frame must begin with the 4-byte magic sequence."""
117 from muse.core.store import _ZSTD_MAGIC, _zstd_compress
118 compressed = _zstd_compress(b"data " * 500)
119 assert compressed[:4] == _ZSTD_MAGIC
120
121 def test_decompress_noop_on_plain_bytes(self) -> None:
122 """Non-zstd bytes are returned unchanged — no corruption."""
123 from muse.core.store import zstd_decompress_if_needed
124 plain = msgpack.packb({"key": "value"}, use_bin_type=True)
125 assert zstd_decompress_if_needed(plain) is plain or zstd_decompress_if_needed(plain) == plain
126
127 def test_decompress_noop_on_empty(self) -> None:
128 from muse.core.store import zstd_decompress_if_needed
129 assert zstd_decompress_if_needed(b"") == b""
130
131 def test_compress_is_smaller_than_input_for_repetitive_data(self) -> None:
132 from muse.core.store import _zstd_compress
133 data = b"aaaa" * 10_000
134 assert len(_zstd_compress(data)) < len(data)
135
136
137 # ---------------------------------------------------------------------------
138 # Tier 2 — Integration: schema_version field in SnapshotRecord
139 # ---------------------------------------------------------------------------
140
141
142 class TestSchemaVersionField:
143 def test_default_schema_version_is_one(self) -> None:
144 """Newly created SnapshotRecord defaults to schema_version=1."""
145 snap = _make_snapshot()
146 assert snap.schema_version == 1
147
148 def test_to_dict_includes_schema_version(self) -> None:
149 """Serialized dict must carry the schema_version key."""
150 snap = _make_snapshot()
151 d = snap.to_dict()
152 assert "schema_version" in d
153 assert d["schema_version"] == 1
154
155 def test_schema_version_excluded_from_snapshot_id_hash(self) -> None:
156 """schema_version is metadata — changing it must not change snapshot_id."""
157 manifest = {"a.py": _obj_id(0xAAAA)}
158 snap_id = compute_snapshot_id(manifest)
159 snap_v1 = SnapshotRecord(snapshot_id=snap_id, manifest=manifest, schema_version=1)
160 snap_v99 = SnapshotRecord(snapshot_id=snap_id, manifest=manifest, schema_version=99)
161 # Both records carry the same snapshot_id; re-verification must pass for both
162 from muse.core.store import _verify_snapshot_id
163 _verify_snapshot_id(snap_v1, snap_id, pathlib.Path("<test>"))
164 _verify_snapshot_id(snap_v99, snap_id, pathlib.Path("<test>"))
165
166 def test_from_msgpack_reads_schema_version(self) -> None:
167 """from_msgpack must deserialise schema_version from the stored dict."""
168 snap = _make_snapshot()
169 d = snap.to_dict()
170 packed = msgpack.packb(d, use_bin_type=True)
171 recovered = SnapshotRecord.from_msgpack(msgpack.unpackb(packed, raw=False))
172 assert recovered.schema_version == 1
173
174 def test_from_msgpack_defaults_schema_version_for_old_files(self) -> None:
175 """Files written before schema_version was added must read as version 1."""
176 snap = _make_snapshot()
177 d = snap.to_dict()
178 del d["schema_version"] # simulate a pre-migration file
179 packed = msgpack.packb(d, use_bin_type=True)
180 recovered = SnapshotRecord.from_msgpack(msgpack.unpackb(packed, raw=False))
181 assert recovered.schema_version == 1
182
183 def test_from_dict_reads_schema_version(self) -> None:
184 snap = _make_snapshot()
185 recovered = SnapshotRecord.from_dict(snap.to_dict())
186 assert recovered.schema_version == 1
187
188 def test_from_dict_defaults_schema_version_for_missing_key(self) -> None:
189 snap = _make_snapshot()
190 d = snap.to_dict()
191 del d["schema_version"]
192 recovered = SnapshotRecord.from_dict(d)
193 assert recovered.schema_version == 1
194
195
196 # ---------------------------------------------------------------------------
197 # Tier 3 — Integration: write / read roundtrip with compression
198 # ---------------------------------------------------------------------------
199
200
201 class TestCompressionRoundtrip:
202 def test_large_snapshot_on_disk_is_zstd_compressed(self, tmp_path: pathlib.Path) -> None:
203 """A manifest large enough to exceed the threshold must be stored as zstd."""
204 from muse.core.store import _ZSTD_COMPRESS_THRESHOLD, _ZSTD_MAGIC
205 root = _init_repo(tmp_path)
206 # Build a snapshot whose packed msgpack exceeds the threshold
207 n = 500 # 500 files; each path is ~30 bytes + sha256 id = generous margin
208 snap = _make_snapshot(n_files=n)
209 write_snapshot(root, snap)
210 path = snapshot_path(root, snap.snapshot_id)
211 raw = path.read_bytes()
212 assert raw[:4] == _ZSTD_MAGIC, (
213 f"Expected zstd magic; got {raw[:4]!r}. "
214 f"Packed size likely below threshold ({_ZSTD_COMPRESS_THRESHOLD} bytes)."
215 )
216
217 def test_small_snapshot_on_disk_is_not_compressed(self, tmp_path: pathlib.Path) -> None:
218 """A tiny manifest below the threshold must be stored as raw msgpack."""
219 from muse.core.store import _ZSTD_MAGIC
220 root = _init_repo(tmp_path)
221 snap = _make_snapshot(n_files=1)
222 write_snapshot(root, snap)
223 raw = snapshot_path(root, snap.snapshot_id).read_bytes()
224 assert raw[:4] != _ZSTD_MAGIC
225
226 def test_compressed_roundtrip_record_is_identical(self, tmp_path: pathlib.Path) -> None:
227 """write_snapshot → read_snapshot must return an identical record (large)."""
228 root = _init_repo(tmp_path)
229 snap = _make_snapshot(n_files=500)
230 write_snapshot(root, snap)
231 loaded = read_snapshot(root, snap.snapshot_id)
232 assert loaded is not None
233 assert loaded.snapshot_id == snap.snapshot_id
234 assert loaded.manifest == snap.manifest
235 assert loaded.directories == snap.directories
236 assert loaded.schema_version == snap.schema_version
237
238 def test_small_roundtrip_record_is_identical(self, tmp_path: pathlib.Path) -> None:
239 """write_snapshot → read_snapshot for a tiny (uncompressed) file."""
240 root = _init_repo(tmp_path)
241 snap = _make_snapshot(n_files=2, note="tiny")
242 write_snapshot(root, snap)
243 loaded = read_snapshot(root, snap.snapshot_id)
244 assert loaded is not None
245 assert loaded.snapshot_id == snap.snapshot_id
246 assert loaded.note == "tiny"
247 assert loaded.schema_version == 1
248
249 def test_schema_version_survives_roundtrip(self, tmp_path: pathlib.Path) -> None:
250 root = _init_repo(tmp_path)
251 snap = _make_snapshot(n_files=500)
252 write_snapshot(root, snap)
253 loaded = read_snapshot(root, snap.snapshot_id)
254 assert loaded is not None
255 assert loaded.schema_version == 1
256
257
258 # ---------------------------------------------------------------------------
259 # Tier 4 — E2E: CLI creates a compressed file on disk
260 # ---------------------------------------------------------------------------
261
262
263 class TestCliCompression:
264 def test_cli_commit_writes_compressed_snapshot(self, tmp_path: pathlib.Path) -> None:
265 """``muse commit`` on a large repo must produce a zstd-compressed snapshot on disk."""
266 from tests.cli_test_helper import CliRunner
267 from muse.core.store import _ZSTD_MAGIC
268 from muse.core.types import fake_id, blob_id
269 from muse.core.object_store import object_path
270
271 runner = CliRunner()
272 env = {"MUSE_REPO_ROOT": str(tmp_path)}
273
274 # Minimal repo structure
275 dot_muse = muse_dir(tmp_path)
276 dot_muse.mkdir()
277 repo_id = fake_id("repo")
278 (dot_muse / "repo.json").write_text(
279 __import__("json").dumps({
280 "repo_id": repo_id, "domain": "code",
281 "default_branch": "main", "created_at": "2025-01-01T00:00:00+00:00",
282 })
283 )
284 (dot_muse / "HEAD").write_text("ref: refs/heads/main")
285 (dot_muse / "refs" / "heads").mkdir(parents=True)
286 for d in ("snapshots", "commits", "objects"):
287 (dot_muse / d).mkdir()
288
289 # Write 300 source files — enough to push past the compression threshold
290 src = tmp_path / "src"
291 src.mkdir()
292 for i in range(300):
293 content = f"module_{i:04d} = {i}\n".encode()
294 obj_id = blob_id(content)
295 obj_path = object_path(tmp_path, obj_id)
296 obj_path.parent.mkdir(parents=True, exist_ok=True)
297 obj_path.write_bytes(content)
298 (src / f"module_{i:04d}.py").write_text(f"module_{i:04d} = {i}\n")
299
300 r = runner.invoke(None, ["commit", "-m", "big"], env=env, catch_exceptions=False)
301 assert r.exit_code == 0, r.output
302
303 snaps_dir = snapshots_dir(tmp_path)
304 snap_files = list(snaps_dir.glob("*/*.msgpack"))
305 assert snap_files, "No snapshot files found after commit"
306 largest = max(snap_files, key=lambda p: p.stat().st_size)
307 raw = largest.read_bytes()
308 assert raw[:4] == _ZSTD_MAGIC, (
309 f"Snapshot {largest.name} not compressed; first 4 bytes: {raw[:4]!r}"
310 )
311
312
313 # ---------------------------------------------------------------------------
314 # Tier 5 — Stress: 1 000-file manifest
315 # ---------------------------------------------------------------------------
316
317
318 class TestStress:
319 def test_1000_file_snapshot_compress_decompress(self, tmp_path: pathlib.Path) -> None:
320 """1 000-file manifest must write and read back correctly under compression."""
321 root = _init_repo(tmp_path)
322 snap = _make_snapshot(n_files=1_000)
323 write_snapshot(root, snap)
324 loaded = read_snapshot(root, snap.snapshot_id)
325 assert loaded is not None
326 assert len(loaded.manifest) == 1_000
327 assert loaded.snapshot_id == snap.snapshot_id
328
329
330 # ---------------------------------------------------------------------------
331 # Tier 6 — State: pre-compression files remain readable
332 # ---------------------------------------------------------------------------
333
334
335 class TestBackwardCompat:
336 def test_old_uncompressed_snapshot_still_readable(self, tmp_path: pathlib.Path) -> None:
337 """A raw-msgpack snapshot written before compression existed must still load."""
338 root = _init_repo(tmp_path)
339 snap = _make_snapshot(n_files=5)
340 path = snapshot_path(root, snap.snapshot_id)
341 path.parent.mkdir(parents=True, exist_ok=True)
342 # Write raw msgpack — no compression, no schema_version
343 raw_dict = snap.to_dict()
344 del raw_dict["schema_version"] # simulate pre-migration file
345 path.write_bytes(msgpack.packb(raw_dict, use_bin_type=True))
346
347 loaded = read_snapshot(root, snap.snapshot_id)
348 assert loaded is not None
349 assert loaded.snapshot_id == snap.snapshot_id
350 assert loaded.schema_version == 1 # default applied
351
352 def test_mixed_compressed_uncompressed_in_same_dir(self, tmp_path: pathlib.Path) -> None:
353 """Both compressed and uncompressed snapshots may coexist in .muse/snapshots/."""
354 root = _init_repo(tmp_path)
355 small = _make_snapshot(n_files=1)
356 large = _make_snapshot(n_files=500)
357
358 write_snapshot(root, small)
359 write_snapshot(root, large)
360
361 loaded_small = read_snapshot(root, small.snapshot_id)
362 loaded_large = read_snapshot(root, large.snapshot_id)
363
364 assert loaded_small is not None
365 assert loaded_large is not None
366 assert len(loaded_large.manifest) == 500
367
368
369 # ---------------------------------------------------------------------------
370 # Tier 7 — Integrity: verify_snapshot_id passes through compression round-trip
371 # ---------------------------------------------------------------------------
372
373
374 class TestIntegrity:
375 def test_verify_snapshot_id_passes_after_compression(self, tmp_path: pathlib.Path) -> None:
376 """Hash verification must succeed when reading a compressed snapshot."""
377 root = _init_repo(tmp_path)
378 snap = _make_snapshot(n_files=500)
379 write_snapshot(root, snap)
380 # read_snapshot internally calls _verify_snapshot_id; None means failure
381 loaded = read_snapshot(root, snap.snapshot_id)
382 assert loaded is not None, "read_snapshot returned None — hash verification failed"
383
384 def test_tampered_compressed_manifest_is_rejected(self, tmp_path: pathlib.Path) -> None:
385 """Altering a byte in a compressed snapshot must cause read_snapshot to return None."""
386 from muse.core.store import _ZSTD_MAGIC
387 root = _init_repo(tmp_path)
388 snap = _make_snapshot(n_files=500)
389 write_snapshot(root, snap)
390 path = snapshot_path(root, snap.snapshot_id)
391 raw = bytearray(path.read_bytes())
392 # Skip the magic and flip a byte deep in the payload
393 if raw[:4] == _ZSTD_MAGIC and len(raw) > 20:
394 raw[16] ^= 0xFF
395 path.write_bytes(bytes(raw))
396 loaded = read_snapshot(root, snap.snapshot_id)
397 assert loaded is None, "Tampered snapshot should not load"
398
399 def test_gc_finds_objects_in_compressed_snapshot(self, tmp_path: pathlib.Path) -> None:
400 """GC reachability walk must extract object IDs from compressed snapshots."""
401 from muse.core.gc import _collect_reachable_objects
402 root = _init_repo(tmp_path)
403 snap = _make_snapshot(n_files=500)
404 write_snapshot(root, snap)
405
406 reachable: set[str] = _collect_reachable_objects(root)
407
408 # All object IDs in the manifest must appear in the reachable set
409 for oid in snap.manifest.values():
410 assert oid in reachable, f"Object {oid[:24]}… not found in GC reachable set"
411
412
413 # ---------------------------------------------------------------------------
414 # Tier 8 — Performance
415 # ---------------------------------------------------------------------------
416
417
418 class TestPerformance:
419 def test_1000_file_roundtrip_under_2s(self, tmp_path: pathlib.Path) -> None:
420 """write_snapshot + read_snapshot for 1 000 files must complete within 2 s."""
421 root = _init_repo(tmp_path)
422 snap = _make_snapshot(n_files=1_000)
423
424 start = time.perf_counter()
425 write_snapshot(root, snap)
426 loaded = read_snapshot(root, snap.snapshot_id)
427 elapsed = time.perf_counter() - start
428
429 assert loaded is not None
430 assert elapsed < 2.0, f"Roundtrip took {elapsed:.3f}s — exceeds 2s budget"
431
432
433 # ---------------------------------------------------------------------------
434 # Tier 9 — Security
435 # ---------------------------------------------------------------------------
436
437
438 class TestSecurity:
439 def test_zstd_bomb_rejected(self, tmp_path: pathlib.Path) -> None:
440 """A zstd-compressed payload that decompresses beyond MAX_MSGPACK_BYTES must be rejected."""
441 import zstandard
442 from muse.core.store import _ZSTD_MAGIC
443
444 root = _init_repo(tmp_path)
445 snap = _make_snapshot(n_files=1)
446
447 # Build a payload that is valid msgpack but huge when decompressed.
448 # We use a repetitive structure that compresses very well.
449 huge_data = b"\x00" * (MAX_MSGPACK_BYTES + 1)
450 compressed = zstandard.ZstdCompressor(level=1).compress(huge_data)
451 assert compressed[:4] == _ZSTD_MAGIC
452
453 path = snapshot_path(root, snap.snapshot_id)
454 path.parent.mkdir(parents=True, exist_ok=True)
455 path.write_bytes(compressed)
456
457 # read_snapshot must fail gracefully — not crash, not return data
458 loaded = read_snapshot(root, snap.snapshot_id)
459 assert loaded is None, "Zstd bomb should be rejected, not loaded"
460
461 def test_schema_version_cannot_alter_snapshot_id(self, tmp_path: pathlib.Path) -> None:
462 """Two records differing only in schema_version must have the same snapshot_id."""
463 manifest = {"src/main.py": _obj_id(0xFFFF)}
464 snap_id = compute_snapshot_id(manifest)
465 r1 = SnapshotRecord(snapshot_id=snap_id, manifest=manifest, schema_version=1)
466 r2 = SnapshotRecord(snapshot_id=snap_id, manifest=manifest, schema_version=42)
467 assert r1.snapshot_id == r2.snapshot_id
468
469 def test_symlinked_snapshot_dir_not_written(self, tmp_path: pathlib.Path) -> None:
470 """write_snapshot must refuse to write when the snapshots/<algo>/ dir is a symlink."""
471 root = _init_repo(tmp_path)
472 snap = _make_snapshot(n_files=1)
473 algo_dir = snapshot_path(root, snap.snapshot_id).parent
474 # Replace the algo dir with a symlink to /tmp
475 if algo_dir.exists():
476 import shutil
477 shutil.rmtree(algo_dir)
478 algo_dir.symlink_to("/tmp")
479 with pytest.raises((ValueError, OSError)):
480 write_snapshot(root, snap)
File History 1 commit
sha256:51ce277f663e01a43eaffbe77509b1de7ac2d4251b55d23306304bcdeb92c90d feat(pack): delta-encode snapshots in MPackBundle wire format Sonnet 4.6 minor ⚠ 121 days ago