gabriel / muse public
test_restore_supercharge.py python
526 lines 19.6 KB
Raw
sha256:51ce277f663e01a43eaffbe77509b1de7ac2d4251b55d23306304bcdeb92c90d feat(pack): delta-encode snapshots in MPackBundle wire format Sonnet 4.6 minor ⚠ breaking 120 days ago
1 """Supercharge tests for ``muse restore`` — performance, data integrity,
2 object-store corruption, concurrency, and source+staged combos.
3
4 Coverage tiers added here:
5 - Performance: duration_ms present, non-negative, and reasonable
6 - Data integrity: complete JSON schema, correct types, exit_code field
7 - Error mapping: object store corruption → exit code 3 (INTERNAL_ERROR)
8 - Concurrent: two threads restore independent files without racing
9 - Source+staged: --source --staged restores stage entry from source commit
10 - Text summary: text output includes "Restored N" summary line
11 - Docstring gap: _resolve_source_manifest returns {} for bad ref (not raises)
12 """
13
14 from __future__ import annotations
15 from collections.abc import Mapping
16
17 import json
18 import pathlib
19 import threading
20 import time
21 import datetime
22 import pytest
23
24 from tests.cli_test_helper import CliRunner, InvokeResult
25
26 from muse.core.object_store import write_object
27 from muse.core.snapshot import compute_commit_id, compute_snapshot_id
28 from muse.core.store import CommitRecord, SnapshotRecord, write_commit, write_snapshot
29 from muse.core.types import Manifest, blob_id
30 from muse.plugins.code.stage import StagedFileMap, make_entry, read_stage, write_stage
31 from muse.core.paths import heads_dir, muse_dir, ref_path
32
33 runner = CliRunner()
34
35 _REPO_ID = "restore-supercharge-test"
36 _counter = 1000 # offset to avoid collisions with test_cmd_restore.py
37
38
39
40
41 def _init_repo(path: pathlib.Path) -> pathlib.Path:
42 muse = muse_dir(path)
43 for d in ("commits", "snapshots", "objects", "refs/heads", "code"):
44 (muse / d).mkdir(parents=True, exist_ok=True)
45 (muse / "HEAD").write_text("ref: refs/heads/main", encoding="utf-8")
46 (muse / "repo.json").write_text(
47 json.dumps({"repo_id": _REPO_ID, "domain": "code"}), encoding="utf-8"
48 )
49 return path
50
51
52 def _env(repo: pathlib.Path) -> Mapping[str, str]:
53 return {"MUSE_REPO_ROOT": str(repo)}
54
55
56 def _commit_files(root: pathlib.Path, files: Mapping[str, bytes], branch: str = "main") -> str:
57 global _counter
58 _counter += 1
59 manifest: Manifest = {}
60 for rel_path, content in files.items():
61 obj_id = blob_id(content)
62 write_object(root, obj_id, content)
63 manifest[rel_path] = obj_id
64 abs_path = root / rel_path
65 abs_path.parent.mkdir(parents=True, exist_ok=True)
66 abs_path.write_bytes(content)
67 snap_id = compute_snapshot_id(manifest)
68 write_snapshot(root, SnapshotRecord(snapshot_id=snap_id, manifest=manifest))
69 committed_at = datetime.datetime.now(datetime.timezone.utc)
70 commit_id = compute_commit_id(
71 parent_ids=[],
72 snapshot_id=snap_id,
73 message=f"commit {_counter}",
74 committed_at_iso=committed_at.isoformat(),
75 )
76 write_commit(
77 root,
78 CommitRecord(
79 commit_id=commit_id,
80 repo_id="test-repo",
81 branch=branch,
82 snapshot_id=snap_id,
83 message=f"commit {_counter}",
84 committed_at=committed_at,
85 ),
86 )
87 (ref_path(root, branch)).write_text(commit_id, encoding="utf-8")
88 return commit_id
89
90
91 def _invoke(repo: pathlib.Path, *args: str) -> InvokeResult:
92 from muse.cli.app import main as cli
93 return runner.invoke(cli, ["restore", *args], env=_env(repo))
94
95
96 # ---------------------------------------------------------------------------
97 # Performance tier
98 # ---------------------------------------------------------------------------
99
100
101 def test_restore_json_has_duration_ms(tmp_path: pathlib.Path) -> None:
102 """JSON output must include 'duration_ms' as a non-negative float."""
103 root = _init_repo(tmp_path)
104 _commit_files(root, {"a.py": b"# orig\n"})
105 (root / "a.py").write_bytes(b"# dirty\n")
106
107 result = _invoke(root, "--json", "a.py")
108 assert result.exit_code == 0
109 data = json.loads(result.stdout)
110 assert "duration_ms" in data, "JSON must include 'duration_ms'"
111 assert isinstance(data["duration_ms"], (int, float)), "duration_ms must be numeric"
112 assert data["duration_ms"] >= 0, "duration_ms must be non-negative"
113
114
115 def test_restore_duration_ms_is_reasonable(tmp_path: pathlib.Path) -> None:
116 """duration_ms for a single-file restore should be well under 5 seconds."""
117 root = _init_repo(tmp_path)
118 _commit_files(root, {"a.py": b"# orig\n"})
119 (root / "a.py").write_bytes(b"# dirty\n")
120
121 result = _invoke(root, "--json", "a.py")
122 assert result.exit_code == 0
123 data = json.loads(result.stdout)
124 assert data["duration_ms"] < 5_000, f"duration_ms={data['duration_ms']} is suspiciously large"
125
126
127 def test_restore_dry_run_json_has_duration_ms(tmp_path: pathlib.Path) -> None:
128 """duration_ms must be present even in dry-run mode."""
129 root = _init_repo(tmp_path)
130 _commit_files(root, {"a.py": b"# orig\n"})
131 (root / "a.py").write_bytes(b"# dirty\n")
132
133 result = _invoke(root, "--dry-run", "--json", "a.py")
134 assert result.exit_code == 0
135 data = json.loads(result.stdout)
136 assert "duration_ms" in data
137
138
139 # ---------------------------------------------------------------------------
140 # Data integrity tier
141 # ---------------------------------------------------------------------------
142
143
144 def test_restore_json_schema_complete_on_success(tmp_path: pathlib.Path) -> None:
145 """All required JSON fields are present with correct types on success."""
146 root = _init_repo(tmp_path)
147 _commit_files(root, {"s.py": b"# orig\n"})
148 (root / "s.py").write_bytes(b"# dirty\n")
149
150 result = _invoke(root, "--json", "s.py")
151 assert result.exit_code == 0
152 data = json.loads(result.stdout)
153
154 assert isinstance(data["restored"], list)
155 assert isinstance(data["not_found"], list)
156 assert isinstance(data["dry_run"], bool)
157 assert isinstance(data["staged"], bool)
158 assert isinstance(data["worktree"], bool)
159 assert isinstance(data["duration_ms"], (int, float))
160 assert isinstance(data["exit_code"], int)
161
162
163 def test_restore_json_exit_code_zero_on_success(tmp_path: pathlib.Path) -> None:
164 """exit_code in JSON is 0 when all files are restored successfully."""
165 root = _init_repo(tmp_path)
166 _commit_files(root, {"ok.py": b"# orig\n"})
167 (root / "ok.py").write_bytes(b"# dirty\n")
168
169 result = _invoke(root, "--json", "ok.py")
170 assert result.exit_code == 0
171 data = json.loads(result.stdout)
172 assert data["exit_code"] == 0
173
174
175 def test_restore_json_exit_code_one_when_file_not_found(tmp_path: pathlib.Path) -> None:
176 """exit_code in JSON is 1 (USER_ERROR) when a file is not in source."""
177 root = _init_repo(tmp_path)
178 _commit_files(root, {"anchor.py": b"# anchor\n"})
179
180 result = _invoke(root, "--json", "ghost.py")
181 assert result.exit_code != 0
182 data = json.loads(result.stdout)
183 assert data["exit_code"] == 1
184
185
186 def test_restore_json_restored_list_correct(tmp_path: pathlib.Path) -> None:
187 """restored list contains exactly the successfully restored paths."""
188 root = _init_repo(tmp_path)
189 _commit_files(root, {"x.py": b"# x\n", "y.py": b"# y\n"})
190 (root / "x.py").write_bytes(b"# dirty x\n")
191 (root / "y.py").write_bytes(b"# dirty y\n")
192
193 result = _invoke(root, "--json", "x.py", "y.py")
194 data = json.loads(result.stdout)
195 assert sorted(data["restored"]) == ["x.py", "y.py"]
196 assert data["not_found"] == []
197
198
199 def test_restore_json_not_found_list_correct(tmp_path: pathlib.Path) -> None:
200 """not_found list contains paths that were absent from the source manifest."""
201 root = _init_repo(tmp_path)
202 _commit_files(root, {"real.py": b"# real\n"})
203 (root / "real.py").write_bytes(b"# dirty\n")
204
205 result = _invoke(root, "--json", "real.py", "ghost.py")
206 data = json.loads(result.stdout)
207 assert "real.py" in data["restored"]
208 assert "ghost.py" in data["not_found"]
209
210
211 def test_restore_json_staged_and_worktree_flags_reflect_args(tmp_path: pathlib.Path) -> None:
212 """staged/worktree fields in JSON reflect the CLI flags used."""
213 root = _init_repo(tmp_path)
214 _commit_files(root, {"f.py": b"# orig\n"})
215 obj_id = blob_id(b"# mod\n")
216 write_object(root, obj_id, b"# mod\n")
217 stage: StagedFileMap = {"f.py": make_entry(obj_id, "M")}
218 write_stage(root, stage)
219
220 result = _invoke(root, "--staged", "--worktree", "--json", "f.py")
221 data = json.loads(result.stdout)
222 assert data["staged"] is True
223 assert data["worktree"] is True
224
225
226 # ---------------------------------------------------------------------------
227 # Error mapping — object store corruption → exit code 3
228 # ---------------------------------------------------------------------------
229
230
231 def test_restore_missing_object_exits_3(tmp_path: pathlib.Path) -> None:
232 """When an object_id is in the manifest but missing from the store, exit code must be 3."""
233 root = _init_repo(tmp_path)
234 content = b"# original\n"
235 obj_id = blob_id(content)
236
237 # Build a manifest pointing at an object that is NOT in the store.
238 # We write the commit but deliberately don't call write_object.
239 manifest: Manifest = {"corrupt.py": obj_id}
240 snap_id = compute_snapshot_id(manifest)
241 write_snapshot(root, SnapshotRecord(snapshot_id=snap_id, manifest=manifest))
242 committed_at = datetime.datetime.now(datetime.timezone.utc)
243 global _counter
244 _counter += 1
245 commit_id = compute_commit_id( parent_ids=[],
246 snapshot_id=snap_id,
247 message=f"corrupt commit {_counter}",
248 committed_at_iso=committed_at.isoformat(),
249 )
250 write_commit(
251 root,
252 CommitRecord(
253 commit_id=commit_id,
254 repo_id="test-repo",
255 branch="main",
256 snapshot_id=snap_id,
257 message=f"corrupt commit {_counter}",
258 committed_at=committed_at,
259 ),
260 )
261 (heads_dir(root) / "main").write_text(commit_id, encoding="utf-8")
262 # Create the file on disk so path resolution doesn't fail
263 (root / "corrupt.py").write_bytes(b"# dirty\n")
264
265 result = _invoke(root, "corrupt.py")
266 assert result.exit_code == 3, (
267 f"Expected exit code 3 (INTERNAL_ERROR) for missing object, got {result.exit_code}"
268 )
269
270
271 def test_restore_missing_object_json_exit_code_3(tmp_path: pathlib.Path) -> None:
272 """JSON exit_code is 3 when the object is missing from the store."""
273 root = _init_repo(tmp_path)
274 content = b"# original\n"
275 obj_id = blob_id(content)
276
277 manifest: Manifest = {"corrupt2.py": obj_id}
278 snap_id = compute_snapshot_id(manifest)
279 write_snapshot(root, SnapshotRecord(snapshot_id=snap_id, manifest=manifest))
280 committed_at = datetime.datetime.now(datetime.timezone.utc)
281 global _counter
282 _counter += 1
283 commit_id = compute_commit_id( parent_ids=[],
284 snapshot_id=snap_id,
285 message=f"corrupt2 {_counter}",
286 committed_at_iso=committed_at.isoformat(),
287 )
288 write_commit(
289 root,
290 CommitRecord(
291 commit_id=commit_id,
292 repo_id="test-repo",
293 branch="main",
294 snapshot_id=snap_id,
295 message=f"corrupt2 {_counter}",
296 committed_at=committed_at,
297 ),
298 )
299 (heads_dir(root) / "main").write_text(commit_id, encoding="utf-8")
300 (root / "corrupt2.py").write_bytes(b"# dirty\n")
301
302 result = _invoke(root, "--json", "corrupt2.py")
303 assert result.exit_code == 3
304 data = json.loads(result.stdout)
305 assert data["exit_code"] == 3
306
307
308 # ---------------------------------------------------------------------------
309 # Concurrent restore
310 # ---------------------------------------------------------------------------
311
312
313 def test_restore_concurrent_independent_files(tmp_path: pathlib.Path) -> None:
314 """Two threads restore independent files without racing or corrupting each other."""
315 root = _init_repo(tmp_path)
316 original_a = b"# thread-a original\n"
317 original_b = b"# thread-b original\n"
318 _commit_files(root, {"ta.py": original_a, "tb.py": original_b})
319 (root / "ta.py").write_bytes(b"# dirty a\n")
320 (root / "tb.py").write_bytes(b"# dirty b\n")
321
322 errors: list[Exception] = []
323
324 def restore_a() -> None:
325 try:
326 result = _invoke(root, "ta.py")
327 assert result.exit_code == 0, f"thread-a exit {result.exit_code}"
328 except Exception as exc:
329 errors.append(exc)
330
331 def restore_b() -> None:
332 try:
333 result = _invoke(root, "tb.py")
334 assert result.exit_code == 0, f"thread-b exit {result.exit_code}"
335 except Exception as exc:
336 errors.append(exc)
337
338 t1 = threading.Thread(target=restore_a)
339 t2 = threading.Thread(target=restore_b)
340 t1.start()
341 t2.start()
342 t1.join(timeout=10)
343 t2.join(timeout=10)
344
345 assert not errors, f"Concurrent restore errors: {errors}"
346 assert (root / "ta.py").read_bytes() == original_a
347 assert (root / "tb.py").read_bytes() == original_b
348
349
350 # ---------------------------------------------------------------------------
351 # --source --staged combo
352 # ---------------------------------------------------------------------------
353
354
355 def test_restore_source_and_staged_clears_stage_from_source(tmp_path: pathlib.Path) -> None:
356 """--source <ref> --staged clears the stage entry so it matches source."""
357 root = _init_repo(tmp_path)
358 v1_content = b"# v1\n"
359 v1_commit = _commit_files(root, {"versioned.py": v1_content})
360
361 # Update to v2
362 v2_content = b"# v2\n"
363 _commit_files(root, {"versioned.py": v2_content})
364
365 # Stage a modification on top of v2
366 mod_content = b"# staged mod\n"
367 obj_id = blob_id(mod_content)
368 write_object(root, obj_id, mod_content)
369 stage: StagedFileMap = {"versioned.py": make_entry(obj_id, "M")}
370 write_stage(root, stage)
371
372 # --source v1_commit --staged should clear the stage entry
373 result = _invoke(root, "--source", v1_commit, "--staged", "versioned.py")
374 assert result.exit_code == 0
375 stage_after = read_stage(root)
376 assert "versioned.py" not in stage_after
377
378
379 def test_restore_source_staged_worktree_restores_from_source(tmp_path: pathlib.Path) -> None:
380 """--source <ref> --staged --worktree restores disk from source, clears stage."""
381 root = _init_repo(tmp_path)
382 v1_content = b"# v1 original\n"
383 v1_commit = _commit_files(root, {"combo.py": v1_content})
384 _commit_files(root, {"combo.py": b"# v2\n"})
385
386 mod_content = b"# staged mod\n"
387 obj_id = blob_id(mod_content)
388 write_object(root, obj_id, mod_content)
389 stage: StagedFileMap = {"combo.py": make_entry(obj_id, "M")}
390 write_stage(root, stage)
391 (root / "combo.py").write_bytes(b"# dirty disk\n")
392
393 result = _invoke(root, "--source", v1_commit, "--staged", "--worktree", "combo.py")
394 assert result.exit_code == 0
395 assert (root / "combo.py").read_bytes() == v1_content
396 stage_after = read_stage(root)
397 assert "combo.py" not in stage_after
398
399
400 # ---------------------------------------------------------------------------
401 # Text summary output
402 # ---------------------------------------------------------------------------
403
404
405 def test_restore_text_output_summary_line(tmp_path: pathlib.Path) -> None:
406 """Text output includes a summary line like 'Restored 2 file(s)'."""
407 root = _init_repo(tmp_path)
408 _commit_files(root, {"p.py": b"# p\n", "q.py": b"# q\n"})
409 (root / "p.py").write_bytes(b"# dirty p\n")
410 (root / "q.py").write_bytes(b"# dirty q\n")
411
412 result = _invoke(root, "p.py", "q.py")
413 assert result.exit_code == 0
414 output = result.stdout + (result.stderr or "")
415 assert "2" in output, f"Expected count in output: {output!r}"
416
417
418 def test_restore_text_output_errors_noted(tmp_path: pathlib.Path) -> None:
419 """Text output notes how many errors occurred when some paths fail."""
420 root = _init_repo(tmp_path)
421 _commit_files(root, {"real.py": b"# real\n"})
422 (root / "real.py").write_bytes(b"# dirty\n")
423
424 result = _invoke(root, "real.py", "ghost.py")
425 assert result.exit_code != 0
426 output = (result.stdout or "") + (result.stderr or "")
427 # Should mention the failure somehow
428 assert "ghost" in output or "error" in output.lower() or "not" in output.lower()
429
430
431 # ---------------------------------------------------------------------------
432 # _resolve_source_manifest — docstring gap: bad ref returns {}, never raises
433 # ---------------------------------------------------------------------------
434
435
436 def test_resolve_source_manifest_bad_ref_returns_empty(tmp_path: pathlib.Path) -> None:
437 """_resolve_source_manifest returns {} for a non-existent ref — never raises."""
438 from muse.cli.commands.restore import _resolve_source_manifest
439 root = _init_repo(tmp_path)
440 _commit_files(root, {"a.py": b"# a\n"})
441 result = _resolve_source_manifest(root, source_ref="nonexistent-branch-xyz")
442 assert result == {}
443
444
445 def test_resolve_source_manifest_valid_ref(tmp_path: pathlib.Path) -> None:
446 """_resolve_source_manifest resolves a valid branch name to its manifest."""
447 from muse.cli.commands.restore import _resolve_source_manifest
448 root = _init_repo(tmp_path)
449 content = b"# branch content\n"
450 _commit_files(root, {"b.py": content}, branch="main")
451 manifest = _resolve_source_manifest(root, source_ref="main")
452 assert "b.py" in manifest
453 assert manifest["b.py"] == blob_id(content)
454
455
456 # ---------------------------------------------------------------------------
457 # Edge: restore staged-only with --source doesn't require file on disk
458 # ---------------------------------------------------------------------------
459
460
461 def test_restore_staged_only_source_does_not_require_disk_file(tmp_path: pathlib.Path) -> None:
462 """--staged with --source works even when the disk file doesn't exist."""
463 root = _init_repo(tmp_path)
464 v1_commit = _commit_files(root, {"staged_only.py": b"# v1\n"})
465 # Stage a modification
466 obj_id = blob_id(b"# mod\n")
467 write_object(root, obj_id, b"# mod\n")
468 stage: StagedFileMap = {"staged_only.py": make_entry(obj_id, "M")}
469 write_stage(root, stage)
470 # Delete disk file
471 (root / "staged_only.py").unlink()
472
473 result = _invoke(root, "--source", v1_commit, "--staged", "staged_only.py")
474 assert result.exit_code == 0
475 stage_after = read_stage(root)
476 assert "staged_only.py" not in stage_after
477
478
479 # ---------------------------------------------------------------------------
480 # Performance: duration_ms for 50-file restore is under 10 seconds
481 # ---------------------------------------------------------------------------
482
483
484 def test_restore_50_files_duration_ms_reasonable(tmp_path: pathlib.Path) -> None:
485 """50-file restore reports duration_ms and completes under 10 seconds."""
486 root = _init_repo(tmp_path)
487 files = {f"perf_{i}.py": f"# orig {i}\n".encode() for i in range(50)}
488 _commit_files(root, files)
489 for name in files:
490 (root / name).write_bytes(b"# dirty\n")
491
492 result = _invoke(root, "--json", *files.keys())
493 assert result.exit_code == 0
494 data = json.loads(result.stdout)
495 assert "duration_ms" in data
496 assert data["duration_ms"] < 10_000
497 assert len(data["restored"]) == 50
498
499
500 class TestRegisterFlags:
501 def test_default_json_out_is_false(self) -> None:
502 import argparse
503 from muse.cli.commands.restore import register
504 p = argparse.ArgumentParser()
505 subs = p.add_subparsers()
506 register(subs)
507 args = p.parse_args(["restore", "src/billing.py"])
508 assert args.json_out is False
509
510 def test_json_flag_sets_json_out(self) -> None:
511 import argparse
512 from muse.cli.commands.restore import register
513 p = argparse.ArgumentParser()
514 subs = p.add_subparsers()
515 register(subs)
516 args = p.parse_args(["restore", "src/billing.py", "--json"])
517 assert args.json_out is True
518
519 def test_j_shorthand_sets_json_out(self) -> None:
520 import argparse
521 from muse.cli.commands.restore import register
522 p = argparse.ArgumentParser()
523 subs = p.add_subparsers()
524 register(subs)
525 args = p.parse_args(["restore", "src/billing.py", "-j"])
526 assert args.json_out is True
File History 1 commit
sha256:51ce277f663e01a43eaffbe77509b1de7ac2d4251b55d23306304bcdeb92c90d feat(pack): delta-encode snapshots in MPackBundle wire format Sonnet 4.6 minor 120 days ago