OBLITERATUS/tests/test_offline_integration.py

"""Offline integration coverage across real model, pipeline, and CLI boundaries."""

from __future__ import annotations

import json
import os
import subprocess
import sys
from pathlib import Path

import pytest
import torch
from datasets import Dataset
from transformers import AutoModelForCausalLM, AutoTokenizer

from obliteratus.abliterate import AbliterationPipeline
from obliteratus.config import DatasetConfig, ModelConfig, StrategyConfig, StudyConfig
from obliteratus.reporting.report import AblationReport
from obliteratus.runner import run_study
from tests.fixtures.tiny_offline_model import build_tiny_offline_model


pytestmark = [pytest.mark.cpu, pytest.mark.integration]


def _state_dict(path: Path) -> dict[str, torch.Tensor]:
    return AutoModelForCausalLM.from_pretrained(
        path,
        local_files_only=True,
    ).state_dict()


def test_fixture_is_deterministic_and_documents_provenance(tmp_path):
    first = build_tiny_offline_model(tmp_path / "first")
    second = build_tiny_offline_model(tmp_path / "second")

    first_state = _state_dict(first)
    second_state = _state_dict(second)
    assert first_state.keys() == second_state.keys()
    assert all(torch.equal(first_state[key], second_state[key]) for key in first_state)

    first_manifest = json.loads((first / "fixture-provenance.json").read_text())
    second_manifest = json.loads((second / "fixture-provenance.json").read_text())
    assert first_manifest == second_manifest
    assert first_manifest["training_data"] is None
    assert first_manifest["third_party_weights"] is None


def test_full_pipeline_saves_and_reloads_a_real_offline_model(tmp_path):
    source = build_tiny_offline_model(tmp_path / "source")
    output = tmp_path / "output"
    events = []
    original = _state_dict(source)

    pipeline = AbliterationPipeline(
        model_name=str(source),
        output_dir=str(output),
        device="cpu",
        dtype="float32",
        method="basic",
        n_directions=1,
        max_seq_length=8,
        verify_sample_size=1,
        harmful_prompts=["harmful request"],
        harmless_prompts=["harmless request"],
        on_stage=events.append,
    )
    result = pipeline.run()

    assert result == output
    assert [(event.stage, event.status) for event in events] == [
        (stage, status)
        for stage in ("summon", "probe", "distill", "excise", "verify", "rebirth")
        for status in ("running", "done")
    ]
    assert (output / "abliteration_metadata.json").is_file()
    assert not list(tmp_path.glob(".output.staging-*"))
    assert not list(tmp_path.glob(".output.backup-*"))

    reloaded = AutoModelForCausalLM.from_pretrained(output, local_files_only=True)
    tokenizer = AutoTokenizer.from_pretrained(output, local_files_only=True)
    batch = tokenizer("hello world", return_tensors="pt")
    with torch.no_grad():
        logits = reloaded(**batch).logits
    assert logits.shape == (1, 2, len(tokenizer))
    assert torch.isfinite(logits).all()
    assert any(
        not torch.equal(original[name], tensor)
        for name, tensor in reloaded.state_dict().items()
    )
    assert set(pipeline._stage_durations) == {
        "summon",
        "probe",
        "distill",
        "excise",
        "verify",
        "rebirth",
    }
    assert all(duration >= 0 for duration in pipeline._stage_durations.values())


def test_installed_wheel_cli_loads_local_model_without_repository_imports(tmp_path):
    source = build_tiny_offline_model(tmp_path / "source")
    isolated_workdir = tmp_path / "outside-repository"
    isolated_workdir.mkdir()
    isolated_home = tmp_path / "home"
    isolated_home.mkdir()
    env = {
        **os.environ,
        "HOME": str(isolated_home),
        "HF_HOME": str(isolated_home / "hf"),
        "HF_DATASETS_OFFLINE": "1",
        "HF_HUB_DISABLE_TELEMETRY": "1",
        "HF_HUB_OFFLINE": "1",
        "TRANSFORMERS_OFFLINE": "1",
    }

    origin = subprocess.run(
        [sys.executable, "-I", "-c", "import obliteratus; print(obliteratus.__file__)"],
        cwd=isolated_workdir,
        env=env,
        check=True,
        capture_output=True,
        text=True,
        timeout=30,
    )
    assert "site-packages" in origin.stdout

    result = subprocess.run(
        [
            sys.executable,
            "-I",
            "-m",
            "obliteratus",
            "info",
            str(source),
            "--device",
            "cpu",
            "--dtype",
            "float32",
        ],
        cwd=isolated_workdir,
        env=env,
        check=True,
        capture_output=True,
        text=True,
        timeout=30,
    )
    assert "architecture:" in result.stdout.lower()
    assert "gpt2" in result.stdout.lower()


def test_study_runner_evaluates_ablates_restores_and_reports(
    tmp_path,
    monkeypatch,
):
    source = build_tiny_offline_model(tmp_path / "source")
    output = tmp_path / "study-results"
    dataset = Dataset.from_dict({"text": ["hello world safe test"]})
    monkeypatch.setattr("obliteratus.runner.load_dataset", lambda **_kwargs: dataset)
    monkeypatch.setattr(AblationReport, "plot_impact", lambda *_args, **_kwargs: None)
    monkeypatch.setattr(AblationReport, "plot_heatmap", lambda *_args, **_kwargs: None)
    config = StudyConfig(
        model=ModelConfig(name=str(source), device="cpu", dtype="float32"),
        dataset=DatasetConfig(name="synthetic/offline", max_samples=1),
        strategies=[StrategyConfig(name="layer_removal")],
        metrics=["perplexity"],
        batch_size=1,
        max_length=8,
        output_dir=str(output),
    )

    report = run_study(config)

    assert report.model_name == str(source)
    assert report.baseline_metrics["perplexity"] > 0
    assert len(report.results) == 1
    assert report.results[0].strategy == "layer_removal"
    assert report.results[0].component == "layer_0"
    assert report.results[0].metrics["perplexity"] > 0
    saved = json.loads((output / "results.json").read_text())
    assert saved["baseline_metrics"] == report.baseline_metrics
    assert (output / "results.csv").is_file()