mirror of
https://github.com/elder-plinius/OBLITERATUS.git
synced 2026-09-21 17:00:50 +02:00
226 lines
8.6 KiB
Python
226 lines
8.6 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
|
|
|
|
|
|
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
|
|
|
|
monkeypatch.delenv("HF_TOKEN", raising=False)
|
|
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
|
|
|
|
|
|
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()
|