OBLITERATUS/tests/test_remaining_cpu_contracts.py

"""Cheap deterministic contracts for remaining CPU-only surfaces."""

from __future__ import annotations

import json
from types import SimpleNamespace
from unittest.mock import Mock

import pytest
import torch

from obliteratus.auto_obliterate import AutoObliterator, IterationResult
from obliteratus.evaluation.advanced_metrics import (
    AbliterationEvalResult,
    _is_degenerate,
    _is_refusal_detailed,
    activation_cosine_similarity,
    effective_rank,
    format_eval_report,
    linear_cka,
    refusal_projection_magnitude,
    refusal_rate_with_ci,
    token_kl_divergence,
)
from obliteratus.evaluation.benchmarks import BenchmarkRunner


class _TinyTokenizer:
    def __call__(self, prompt, return_tensors="pt", truncation=True, max_length=256):
        return {"input_ids": torch.tensor([[1, 2, 3]])}

    def encode(self, letter, add_special_tokens=False):
        return [ord(letter)]

    def decode(self, tokens, skip_special_tokens=True):
        return ""


class _TinyModel:
    def __init__(self):
        self._p = torch.nn.Parameter(torch.zeros(1))

    def parameters(self):
        return iter([self._p])

    def __call__(self, **_inputs):
        return SimpleNamespace(logits=torch.tensor([[[0.1, 0.9, 0.0, -0.1]]]))

    def generate(self, **_inputs):
        return torch.tensor([[1, 2, 3, 4]])


def test_auto_obliterate_resume_loads_valid_state_and_saves_round_trip(tmp_path, monkeypatch):
    monkeypatch.setattr("obliteratus.auto_obliterate.Path.home", staticmethod(lambda: tmp_path))
    output_base = tmp_path / "auto"
    state_file = output_base / "auto_state.json"
    state_file.parent.mkdir(parents=True)
    state_file.write_text(
        json.dumps(
            {
                "model_id": "demo",
                "iterations": [
                    {
                        "iteration": 1,
                        "method": "aggressive",
                        "prompt_volume": 4,
                        "time_seconds": 1.0,
                        "output_dir": str(output_base / "iter_1"),
                    }
                ],
            }
        )
    )

    ob = AutoObliterator("demo", max_iterations=2, output_base=str(output_base))
    assert ob._resume_from == 1
    assert len(ob._result.iterations) == 1

    ob._result.iterations.append(
        IterationResult(iteration=2, method="nuclear", prompt_volume=8, output_dir=str(output_base / "iter_2"))
    )
    ob._save_state()
    saved = json.loads(state_file.read_text())
    assert saved["iterations"][0]["iteration"] == 1
    assert saved["iterations"][1]["method"] == "nuclear"


def test_auto_obliterate_quarantines_invalid_state_and_prompt_fallback(monkeypatch, tmp_path):
    monkeypatch.setattr("obliteratus.auto_obliterate.Path.home", staticmethod(lambda: tmp_path))
    output_base = tmp_path / "auto"
    state_file = output_base / "auto_state.json"
    state_file.parent.mkdir(parents=True)
    state_file.write_text("{not-json")

    ob = AutoObliterator("demo", output_base=str(output_base))
    quarantined = list(output_base.glob("auto_state.json.corrupt-*"))
    assert quarantined
    assert ob._resume_from == 0

    monkeypatch.setitem(
        __import__("sys").modules,
        "obliteratus.prompts",
        SimpleNamespace(
            load_dataset_source=Mock(side_effect=RuntimeError("boom")),
            DEFAULT_HARMFUL=("h1",),
            DEFAULT_HARMLESS=("n1",),
        ),
    )
    harmful, harmless = AutoObliterator._get_expanded_prompts(0)
    assert harmful == ["h1"]
    assert harmless == ["n1"]


def test_benchmark_runner_generation_truthfulness_and_number_branches():
    runner = BenchmarkRunner(_TinyModel(), _TinyTokenizer(), device="cpu")
    assert runner._generate_short("prompt") == ""
    assert runner._check_truthfulness("No, that is a myth.", "no", "yes") is True
    assert runner._check_truthfulness("yes", "no", "yes") is False
    assert runner._extract_number("The answer is 42.") == 42.0
    assert runner._extract_number("no number here") is None


def test_advanced_metrics_tensor_validation_refusal_and_report_edges():
    with pytest.raises(ValueError, match="identical shapes"):
        token_kl_divergence(torch.zeros(1, 2, 3), torch.zeros(1, 3, 3))

    with pytest.raises(ValueError, match="must not be empty"):
        token_kl_divergence(torch.zeros(1, 0, 3), torch.zeros(1, 0, 3))

    with pytest.raises(ValueError, match="same sample count"):
        linear_cka(torch.zeros(1, 2), torch.zeros(2, 2))

    with pytest.raises(ValueError, match="non-zero norm"):
        refusal_projection_magnitude(torch.ones(2, 3), torch.zeros(3))

    assert _is_degenerate("!!!!!")
    assert _is_refusal_detailed("I cannot help with that.", "combined")[0] is True
    assert _is_refusal_detailed("   ", "combined") == (False, "")

    report = format_eval_report(
        AbliterationEvalResult(
            refusal_rate_harmful=0.0,
            refusal_rate_harmless=None,
            kl_divergence=0.1,
            perplexity=None,
            coherence_score=0.2,
            mean_activation_cosine=None,
            mean_cka=0.3,
        )
    )
    assert "unavailable" in report
    assert "excellent" in report

    assert refusal_rate_with_ci([], mode="combined") == {
        "available": False,
        "rate": None,
        "ci_lower": None,
        "ci_upper": None,
        "n_samples": 0,
        "refusal_count": 0,
    }

    assert effective_rank(torch.eye(2)) == pytest.approx(2.0)
    assert activation_cosine_similarity(torch.ones(2, 3), torch.ones(2, 3)) == pytest.approx(1.0)