"""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