OBLITERATUS/tests/test_interactive_contracts.py

"""Scripted contracts for the guided interactive workflow."""

from __future__ import annotations

from types import SimpleNamespace
from unittest.mock import Mock

import pytest

from obliteratus import interactive
from obliteratus.presets import ModelPreset


def _preset(*, quantization: str | None = None) -> ModelPreset:
    return ModelPreset(
        name="Safe model",
        hf_id="org/safe-model",
        description="fixture",
        tier="tiny",
        params="1B",
        recommended_dtype="float32",
        recommended_quantization=quantization,
    )


@pytest.mark.parametrize(
    ("vram_gb", "expected"),
    [(4, "small"), (8, "medium"), (19.9, "medium"), (20, "large")],
)
def test_compute_tier_maps_cuda_memory_boundaries(monkeypatch, vram_gb, expected):
    from obliteratus import device

    monkeypatch.setattr(device, "is_cuda", lambda: True)
    monkeypatch.setattr(device, "is_mps", lambda: False)
    torch = SimpleNamespace(
        cuda=SimpleNamespace(
            get_device_properties=lambda _index: SimpleNamespace(
                total_memory=vram_gb * 1024**3,
            )
        )
    )
    monkeypatch.setitem(__import__("sys").modules, "torch", torch)

    assert interactive._detect_compute_tier() == expected


@pytest.mark.parametrize(
    ("memory_gb", "expected"),
    [(16, "small"), (24, "medium")],
)
def test_compute_tier_maps_mps_unified_memory(monkeypatch, memory_gb, expected):
    from obliteratus import device

    monkeypatch.setattr(device, "is_cuda", lambda: False)
    monkeypatch.setattr(device, "is_mps", lambda: True)
    monkeypatch.setattr(
        device,
        "get_memory_info",
        lambda: SimpleNamespace(total_gb=memory_gb),
    )

    assert interactive._detect_compute_tier() == expected


def test_compute_tier_falls_back_to_cpu_when_device_probe_is_unavailable(monkeypatch):
    from obliteratus import device

    monkeypatch.setattr(device, "is_cuda", Mock(side_effect=ImportError("torch")))

    assert interactive._detect_compute_tier() == "tiny"


def test_custom_model_selection_preserves_safe_tier_defaults(monkeypatch):
    monkeypatch.setattr(interactive, "get_presets_by_tier", lambda _tier: [_preset()])
    monkeypatch.setattr(interactive.IntPrompt, "ask", lambda *_args, **_kwargs: 0)
    monkeypatch.setattr(
        interactive.Prompt,
        "ask",
        lambda *_args, **_kwargs: "org/custom-model",
    )

    selected = interactive._pick_model("tiny")

    assert selected.hf_id == "org/custom-model"
    assert selected.recommended_dtype == "float32"
    assert selected.recommended_quantization is None


def test_invalid_model_selection_falls_back_to_first_recommendation(monkeypatch):
    first = _preset()
    monkeypatch.setattr(interactive, "get_presets_by_tier", lambda _tier: [first])
    monkeypatch.setattr(interactive.IntPrompt, "ask", lambda *_args, **_kwargs: 99)

    assert interactive._pick_model("tiny") is first


def test_custom_strategy_and_sample_mappings_are_exact(monkeypatch):
    answers = iter(["5", "3"])
    monkeypatch.setattr(
        interactive.Prompt,
        "ask",
        lambda *_args, **_kwargs: next(answers),
    )

    strategies = interactive._pick_strategies()

    assert [item["name"] for item in strategies] == [
        "layer_removal",
        "head_pruning",
        "ffn_ablation",
        "embedding_ablation",
    ]
    assert strategies[-1]["params"] == {"chunk_size": 48}
    assert interactive._pick_sample_size() == 500


def test_guided_run_builds_safe_config_and_returns_study_result(monkeypatch, tmp_path):
    from obliteratus import device, runner

    preset = _preset()
    study = SimpleNamespace(
        name="Fast",
        strategies=[{"name": "layer_removal", "params": {"limit": 1}}],
        max_samples=7,
        batch_size=2,
        max_length=64,
    )
    monkeypatch.chdir(tmp_path)
    monkeypatch.setattr(interactive, "_pick_compute_tier", lambda: "tiny")
    monkeypatch.setattr(interactive, "_pick_model", lambda _tier: preset)
    monkeypatch.setattr(interactive, "_pick_study_preset", lambda: study)
    monkeypatch.setattr(device, "get_device", lambda: "cpu")
    monkeypatch.setattr(interactive.Confirm, "ask", lambda *_args, **_kwargs: True)
    run_study = Mock(return_value="study-result")
    monkeypatch.setattr(runner, "run_study", run_study)

    assert interactive.run_interactive() == "study-result"
    config = run_study.call_args.args[0]
    assert config.model.name == "org/safe-model"
    assert config.model.device == "cpu"
    assert config.model.trust_remote_code is False
    assert config.dataset.max_samples == 7
    assert config.output_dir == "results/org_safe-model"


def test_guided_run_cancellation_does_not_start_a_study(monkeypatch):
    from obliteratus import device, runner

    monkeypatch.setattr(interactive, "_pick_compute_tier", lambda: "tiny")
    monkeypatch.setattr(interactive, "_pick_model", lambda _tier: _preset())
    monkeypatch.setattr(
        interactive,
        "_pick_study_preset",
        lambda: SimpleNamespace(
            name="Fast",
            strategies=[{"name": "layer_removal", "params": {}}],
            max_samples=1,
            batch_size=1,
            max_length=8,
        ),
    )
    monkeypatch.setattr(device, "get_device", lambda: "cpu")
    monkeypatch.setattr(interactive.Confirm, "ask", lambda *_args, **_kwargs: False)
    run_study = Mock()
    monkeypatch.setattr(runner, "run_study", run_study)

    assert interactive.run_interactive() is None
    run_study.assert_not_called()


def test_quantized_guided_run_returns_quantized_result(monkeypatch):
    from obliteratus import device

    monkeypatch.setattr(interactive, "_pick_compute_tier", lambda: "small")
    monkeypatch.setattr(
        interactive,
        "_pick_model",
        lambda _tier: _preset(quantization="4bit"),
    )
    monkeypatch.setattr(
        interactive,
        "_pick_study_preset",
        lambda: SimpleNamespace(
            name="Fast",
            strategies=[{"name": "layer_removal", "params": {}}],
            max_samples=1,
            batch_size=1,
            max_length=8,
        ),
    )
    monkeypatch.setattr(device, "get_device", lambda: "cuda")
    monkeypatch.setattr(interactive.Confirm, "ask", lambda *_args, **_kwargs: True)
    run_quantized = Mock(return_value="quantized-result")
    monkeypatch.setattr(interactive, "_run_quantized", run_quantized)

    assert interactive.run_interactive() == "quantized-result"
    config, quantization = run_quantized.call_args.args
    assert quantization == "4bit"
    assert config.model.device == "auto"


def test_quantized_runner_sets_loader_contract_before_execution(monkeypatch):
    from obliteratus import runner

    config = SimpleNamespace(model=SimpleNamespace(device="cuda", quantization=None))
    run_study = Mock(return_value="done")
    monkeypatch.setattr(runner, "run_study", run_study)

    assert interactive._run_quantized(config, "8bit") == "done"
    assert config.model.device == "auto"
    assert config.model.quantization == "8bit"
    run_study.assert_called_once_with(config)