test: cover notebook stage callback contract (#30)

tests/test_notebook_contract.py

@@ -0,0 +1,39 @@
+"""Regression tests for executable contracts embedded in project notebooks."""
+
+from __future__ import annotations
+
+import ast
+import json
+from pathlib import Path
+from types import SimpleNamespace
+
+
+def test_abliterate_notebook_stage_callback_uses_stage_result_contract(capsys):
+    """The Colab callback should consume StageResult.stage and .message."""
+    notebook = json.loads(Path("notebooks/abliterate.ipynb").read_text())
+    callback = None
+
+    for cell in notebook["cells"]:
+        if cell.get("cell_type") != "code":
+            continue
+        source = "".join(cell.get("source", []))
+        if "def on_stage" not in source:
+            continue
+        tree = ast.parse(source)
+        for node in tree.body:
+            if isinstance(node, ast.FunctionDef) and node.name == "on_stage":
+                callback = node
+                break
+        if callback is not None:
+            break
+
+    assert callback is not None, "notebook no longer defines on_stage"
+    namespace: dict[str, object] = {}
+    ast.fix_missing_locations(callback)
+    exec(compile(ast.Module(body=[callback], type_ignores=[]), "<notebook>", "exec"), namespace)
+
+    namespace["on_stage"](SimpleNamespace(stage="probe", message="loading model"))
+    output = capsys.readouterr().out
+
+    assert "STAGE: PROBE" in output
+    assert "loading model" in output