From c2f903491473eb8f9e50f2b14ec4589fe60f01f1 Mon Sep 17 00:00:00 2001 From: Alex Cheema Date: Fri, 13 Feb 2026 10:40:54 -0800 Subject: [PATCH] feat: add reusable snapshot regression testing to e2e framework Add e2e/snapshot.py with assert_snapshot() for deterministic regression testing. On first run, saves inference output as the expected snapshot. On subsequent runs, compares against it with unified diff on mismatch. Set UPDATE_SNAPSHOTS=1 or pass --update-snapshots to regenerate. Refactor test_inference_snapshot.py to use the shared infrastructure and drop temperature=0 in favor of seed-only determinism. Co-Authored-By: Claude Opus 4.6 --- e2e/run_all.py | 2 ++ e2e/snapshot.py | 64 ++++++++++++++++++++++++++++++++++ e2e/snapshots/inference.json | 8 ----- e2e/test_inference_snapshot.py | 48 ++++++++----------------- 4 files changed, 80 insertions(+), 42 deletions(-) create mode 100644 e2e/snapshot.py delete mode 100644 e2e/snapshots/inference.json diff --git a/e2e/run_all.py b/e2e/run_all.py index 0ab2aafa..171edde0 100644 --- a/e2e/run_all.py +++ b/e2e/run_all.py @@ -34,6 +34,8 @@ def is_slow(test_file: Path) -> bool: def main(): run_slow = "--slow" in sys.argv or os.environ.get("E2E_SLOW") == "1" + if "--update-snapshots" in sys.argv: + os.environ["UPDATE_SNAPSHOTS"] = "1" test_files = sorted(E2E_DIR.glob("test_*.py")) if not test_files: print("No test files found") diff --git a/e2e/snapshot.py b/e2e/snapshot.py new file mode 100644 index 00000000..3a03e33b --- /dev/null +++ b/e2e/snapshot.py @@ -0,0 +1,64 @@ +"""Snapshot testing infrastructure for E2E tests. + +Provides deterministic regression testing by comparing inference output +against saved snapshots. On first run, snapshots are created automatically. +Set UPDATE_SNAPSHOTS=1 to regenerate snapshots when output intentionally changes. +""" + +import difflib +import json +import os +from pathlib import Path + +SNAPSHOTS_DIR = Path(__file__).parent / "snapshots" + + +def assert_snapshot( + name: str, + content: str, + metadata: dict, +) -> None: + """Compare content against a saved snapshot, or create one if missing. + + Args: + name: Snapshot identifier (used as filename: snapshots/{name}.json). + content: The actual inference output to compare. + metadata: Additional context stored alongside content (model, seed, etc.). + Not used for comparison -- purely documentary. + + Raises: + AssertionError: If content doesn't match the saved snapshot. + + Environment: + UPDATE_SNAPSHOTS=1: Overwrite existing snapshot with actual content. + """ + snapshot_file = SNAPSHOTS_DIR / f"{name}.json" + update = os.environ.get("UPDATE_SNAPSHOTS") == "1" + + if snapshot_file.exists() and not update: + snapshot = json.loads(snapshot_file.read_text()) + expected = snapshot["content"] + if content != expected: + diff = "\n".join( + difflib.unified_diff( + expected.splitlines(), + content.splitlines(), + fromfile=f"expected ({snapshot_file.relative_to(SNAPSHOTS_DIR.parent.parent)})", + tofile="actual", + lineterm="", + ) + ) + raise AssertionError( + f"Snapshot mismatch for '{name}'!\n\n" + f"{diff}\n\n" + f"Expected: {expected!r}\n" + f"Actual: {content!r}\n\n" + f"To update: UPDATE_SNAPSHOTS=1 python3 e2e/run_all.py --slow" + ) + print(f" Output matches snapshot ({snapshot_file.name})") + else: + SNAPSHOTS_DIR.mkdir(parents=True, exist_ok=True) + snapshot_data = {**metadata, "content": content} + snapshot_file.write_text(json.dumps(snapshot_data, indent=2) + "\n") + action = "Updated" if update else "Created" + print(f" {action} snapshot: {snapshot_file}") diff --git a/e2e/snapshots/inference.json b/e2e/snapshots/inference.json deleted file mode 100644 index d2721114..00000000 --- a/e2e/snapshots/inference.json +++ /dev/null @@ -1,8 +0,0 @@ -{ - "model": "mlx-community/Qwen3-0.6B-4bit", - "seed": 42, - "temperature": 0, - "prompt": "What is 2+2? Reply with just the number.", - "max_tokens": 32, - "content": "\nOkay, so I need to figure out what 2+2 is. Let me think. Well, if you add 2 and 2 together" -} diff --git a/e2e/test_inference_snapshot.py b/e2e/test_inference_snapshot.py index 79d88e8f..03f522ba 100644 --- a/e2e/test_inference_snapshot.py +++ b/e2e/test_inference_snapshot.py @@ -1,7 +1,7 @@ """Test: Deterministic inference output (snapshot test). slow -Sends a chat completion request with a fixed seed and temperature=0, +Sends a chat completion request with a fixed seed, then verifies the output matches a known-good snapshot. This ensures inference produces consistent results across runs. @@ -10,18 +10,18 @@ Run with: python3 e2e/run_all.py --slow or E2E_SLOW=1 python3 e2e/run_all.py """ import asyncio -import json import sys from pathlib import Path sys.path.insert(0, str(Path(__file__).parent)) +from snapshot import assert_snapshot + from conftest import Cluster MODEL = "mlx-community/Qwen3-0.6B-4bit" SEED = 42 PROMPT = "What is 2+2? Reply with just the number." MAX_TOKENS = 32 -SNAPSHOT_FILE = Path(__file__).parent / "snapshots" / "inference.json" async def main(): @@ -30,50 +30,30 @@ async def main(): await cluster.start() await cluster.assert_healthy() - # Launch the model instance (triggers download + placement) print(f" Launching model {MODEL}...") await cluster.place_model(MODEL) - print(f" Sending chat completion (seed={SEED}, temperature=0)...") + print(f" Sending chat completion (seed={SEED})...") resp = await cluster.chat( model=MODEL, messages=[{"role": "user", "content": PROMPT}], seed=SEED, - temperature=0, max_tokens=MAX_TOKENS, ) content = resp["choices"][0]["message"]["content"] print(f" Response: {content!r}") - # Load or create snapshot - if SNAPSHOT_FILE.exists(): - snapshot = json.loads(SNAPSHOT_FILE.read_text()) - expected = snapshot["content"] - assert content == expected, ( - f"Snapshot mismatch!\n" - f" Expected: {expected!r}\n" - f" Got: {content!r}\n" - f" Delete {SNAPSHOT_FILE} to regenerate." - ) - print(" Output matches snapshot") - else: - SNAPSHOT_FILE.parent.mkdir(parents=True, exist_ok=True) - SNAPSHOT_FILE.write_text( - json.dumps( - { - "model": MODEL, - "seed": SEED, - "temperature": 0, - "prompt": PROMPT, - "max_tokens": MAX_TOKENS, - "content": content, - }, - indent=2, - ) - + "\n" - ) - print(f" Snapshot created: {SNAPSHOT_FILE}") + assert_snapshot( + name="inference_snapshot", + content=content, + metadata={ + "model": MODEL, + "seed": SEED, + "prompt": PROMPT, + "max_tokens": MAX_TOKENS, + }, + ) print("PASSED: inference_snapshot")