From 9e546f7bc2d0eca2378b6c3b22847a42389fab81 Mon Sep 17 00:00:00 2001
From: Joseph Magly <1159087+jmagly@users.noreply.github.com>
Date: Fri, 14 Aug 2026 10:30:18 -0400
Subject: [PATCH] test: cover notebook stage callback contract (#30)

---
 tests/test_notebook_contract.py | 39 +++++++++++++++++++++++++++++++++
 1 file changed, 39 insertions(+)
 create mode 100644 tests/test_notebook_contract.py

diff --git a/tests/test_notebook_contract.py b/tests/test_notebook_contract.py
new file mode 100644
index 0000000..6080b30
--- /dev/null
+++ b/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