mirror of
https://github.com/elder-plinius/OBLITERATUS.git
synced 2026-08-17 16:37:30 +02:00
40 lines
1.3 KiB
Python
40 lines
1.3 KiB
Python
"""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
|