OBLITERATUS/tests/test_notebook_contract.py

"""Regression tests for executable contracts embedded in project notebooks."""

from __future__ import annotations

import ast
import json
import os
from pathlib import Path
from types import SimpleNamespace

import pytest


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


def _notebook_code(marker):
    notebook = json.loads(Path("notebooks/abliterate.ipynb").read_text())
    return next(
        "".join(cell["source"])
        for cell in notebook["cells"]
        if cell["cell_type"] == "code" and marker in "".join(cell["source"])
    )


def _run_harness(monkeypatch, *, token=None, failure=None):
    """Execute real notebook cells with an offline Hub and observable pipeline."""
    import sys
    from unittest.mock import Mock

    # Record even an initially absent key before notebook code writes directly
    # to os.environ, so teardown restores the caller's original environment.
    monkeypatch.setenv("HF_TOKEN", "")
    monkeypatch.delenv("HF_TOKEN")
    events = []

    def download(**kwargs):
        events.append(("access", kwargs))
        if failure:
            raise failure
        return "offline-config.json"

    def pipeline(**kwargs):
        import os
        events.append(("construct", kwargs, os.environ.get("HF_TOKEN")))
        return SimpleNamespace(run=lambda: events.append(("run",)) or "output")

    hub = SimpleNamespace(
        get_token=Mock(return_value=token), hf_hub_download=download,
        HfApi=Mock(side_effect=AssertionError("Unexpected upload")),
        notebook_login=Mock(side_effect=AssertionError("Unexpected login widget")),
    )
    secret = Mock(return_value=None)
    monkeypatch.setitem(sys.modules, "huggingface_hub", hub)
    monkeypatch.setitem(sys.modules, "obliteratus.credential_sources", SimpleNamespace(resolve_secret=secret))
    monkeypatch.setitem(sys.modules, "obliteratus.abliterate", SimpleNamespace(AbliterationPipeline=pipeline))
    namespace = {}
    exec(_notebook_code("#@title Abliteration Config"), namespace)
    return namespace, events, hub, secret


@pytest.mark.parametrize("original_token", [None, "fake-caller-token"])
@pytest.mark.parametrize("access_denied", [False, True])
def test_notebook_harness_restores_token_after_execution(monkeypatch, original_token, access_denied):
    if original_token is None:
        monkeypatch.delenv("HF_TOKEN", raising=False)
    else:
        monkeypatch.setenv("HF_TOKEN", original_token)

    with monkeypatch.context() as notebook_patch:
        namespace, _, _, _ = _run_harness(
            notebook_patch,
            token="fake-notebook-token",
            failure=PermissionError("denied") if access_denied else None,
        )
        source = _notebook_code("def check_model_access")
        if access_denied:
            with pytest.raises(RuntimeError, match="Model access"):
                exec(source, namespace)
        else:
            exec(source, namespace)
            assert os.environ["HF_TOKEN"] == "fake-notebook-token"

    assert os.environ.get("HF_TOKEN") == original_token


def test_default_gated_model_stops_before_pipeline_without_access(monkeypatch, capsys):
    import traceback
    import pytest

    namespace, events, hub, _ = _run_harness(
        monkeypatch, failure=PermissionError("private request details fake-secret-value"),
    )
    with pytest.raises(RuntimeError, match="HF_TOKEN") as caught:
        exec(_notebook_code("def check_model_access"), namespace)
    assert [event[0] for event in events] == ["access"]
    assert events[0][1] == {
        "repo_id": "meta-llama/Llama-3.1-8B-Instruct", "filename": "config.json",
        "token": False, "force_download": True,
    }
    rendered = "".join(traceback.format_exception(caught.type, caught.value, caught.tb))
    assert "fake-secret-value" not in rendered + capsys.readouterr().out
    assert "wait for approval" in str(caught.value)
    hub.notebook_login.assert_not_called()


def test_authorized_colab_token_reaches_loader_without_output(monkeypatch, capsys):
    namespace, events, _, _ = _run_harness(monkeypatch, token="fake-private-read-token")
    exec(_notebook_code("def check_model_access"), namespace)
    assert [event[0] for event in events] == ["access", "construct", "run"]
    assert events[0][1]["token"] == "fake-private-read-token"
    assert events[1][2] == "fake-private-read-token"
    assert "fake-private-read-token" not in capsys.readouterr().out


def test_ungated_model_runs_anonymously_and_upload_is_opt_in(monkeypatch):
    namespace, events, hub, _ = _run_harness(monkeypatch)
    namespace["MODEL"] = "Qwen/Qwen2.5-7B-Instruct"
    exec(_notebook_code("def check_model_access"), namespace)
    exec(_notebook_code("UPLOAD_TO_HUB ="), namespace)
    assert [event[0] for event in events] == ["access", "construct", "run"]
    assert events[0][1]["token"] is False
    assert events[1][2] is None
    hub.HfApi.assert_not_called()
    hub.notebook_login.assert_not_called()


def test_direct_rerun_rechecks_changed_model_and_blocks_expired_access(monkeypatch):
    import pytest

    namespace, events, hub, _ = _run_harness(monkeypatch, token="fake-token")
    source = _notebook_code("def check_model_access")
    exec(source, namespace)
    namespace["MODEL"] = "other/private-model"

    def denied(**kwargs):
        events.append(("access", kwargs))
        raise PermissionError("expired")

    hub.hf_hub_download = denied
    with pytest.raises(RuntimeError, match="Model access"):
        exec(source, namespace)
    assert [event[0] for event in events] == ["access", "construct", "run", "access"]
    assert events[-1][1]["repo_id"] == "other/private-model"


def test_configured_secret_precedes_hub_cached_token(monkeypatch):
    namespace, events, hub, secret = _run_harness(monkeypatch, token="fake-other-account")
    secret.return_value = "fake-configured-token"
    exec(_notebook_code("def check_model_access"), namespace)
    assert events[0][1]["token"] == "fake-configured-token"
    hub.get_token.assert_not_called()


def test_secret_resolution_failure_cannot_fall_back_or_run(monkeypatch):
    import pytest

    namespace, events, hub, secret = _run_harness(monkeypatch, token="fake-cached-token")
    secret.side_effect = RuntimeError("sensitive credential source")
    with pytest.raises(RuntimeError, match="Model access"):
        exec(_notebook_code("def check_model_access"), namespace)
    assert events == []
    hub.get_token.assert_not_called()


def test_notebook_preflight_network_failure_stops_before_loading(monkeypatch):
    import pytest

    namespace, events, _, _ = _run_harness(monkeypatch, failure=ConnectionError("offline"))
    with pytest.raises(RuntimeError, match="connection"):
        exec(_notebook_code("def check_model_access"), namespace)
    assert [event[0] for event in events] == ["access"]


def test_notebook_has_no_persisted_execution_outputs():
    notebook = json.loads(Path("notebooks/abliterate.ipynb").read_text())
    for cell in notebook["cells"]:
        if cell["cell_type"] == "code":
            assert cell.get("outputs", []) == []
            assert cell.get("execution_count") is None


def test_token_without_gated_approval_cannot_run(monkeypatch):
    import pytest

    namespace, events, _, _ = _run_harness(
        monkeypatch, token="fake-unapproved-token", failure=PermissionError("403"),
    )
    with pytest.raises(RuntimeError, match="wait for approval"):
        exec(_notebook_code("def check_model_access"), namespace)
    assert [event[0] for event in events] == ["access"]


def test_explicit_upload_uses_resolved_token_and_destination(monkeypatch):
    from unittest.mock import Mock

    namespace, _, hub, secret = _run_harness(monkeypatch)
    secret.return_value = "fake-write-token"
    api = Mock()
    hub.HfApi = Mock(return_value=api)
    namespace["model_dir"] = Path("saved-model")
    source = _notebook_code("UPLOAD_TO_HUB =").replace(
        "UPLOAD_TO_HUB = False", "UPLOAD_TO_HUB = True",
    ).replace(
        'HF_REPO = "your-username/model-name-abliterated"',
        'HF_REPO = "researcher/output"',
    )
    exec(source, namespace)
    hub.HfApi.assert_called_once_with(token="fake-write-token")
    api.create_repo.assert_called_once_with("researcher/output", exist_ok=True)
    api.upload_folder.assert_called_once_with(
        folder_path="saved-model", repo_id="researcher/output", repo_type="model",
    )


def test_explicit_upload_stops_without_credentials(monkeypatch):
    import pytest

    namespace, _, hub, _ = _run_harness(monkeypatch)
    source = _notebook_code("UPLOAD_TO_HUB =").replace(
        "UPLOAD_TO_HUB = False", "UPLOAD_TO_HUB = True",
    ).replace(
        'HF_REPO = "your-username/model-name-abliterated"',
        'HF_REPO = "researcher/output"',
    )
    with pytest.raises(RuntimeError, match="write token"):
        exec(source, namespace)
    hub.HfApi.assert_not_called()