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=[]), "", "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