diff --git a/tests/test_capability_check.py b/tests/test_capability_check.py new file mode 100644 index 0000000..3e13737 --- /dev/null +++ b/tests/test_capability_check.py @@ -0,0 +1,106 @@ +"""Tests for obliteratus.capability_check.""" + +from __future__ import annotations + +import json +from pathlib import Path +from unittest.mock import MagicMock, patch + +import pytest + + +class TestCapabilityCheckImport: + def test_import(self): + from obliteratus.capability_check import capability_check # noqa: F401 + + def test_quick_subjects_defined(self): + from obliteratus.capability_check import QUICK_SUBJECTS + + assert len(QUICK_SUBJECTS) >= 3 + assert all(s.startswith("mmlu_") for s in QUICK_SUBJECTS) + + +class TestRunLmEval: + @patch("obliteratus.capability_check.subprocess.run") + def test_calls_lm_eval(self, mock_run, tmp_path): + from obliteratus.capability_check import _run_lm_eval + + # Create fake results + results_dir = tmp_path / "results" / "fake_model" + results_dir.mkdir(parents=True) + results_file = results_dir / "results_2026.json" + results_file.write_text(json.dumps({ + "results": {"mmlu": {"acc,none": 0.85}}, + })) + + mock_run.return_value = MagicMock(returncode=0, stdout="", stderr="") + + result = _run_lm_eval("fake/model", "mmlu", 5, "cpu", str(tmp_path / "results")) + assert result["results"]["mmlu"]["acc,none"] == 0.85 + mock_run.assert_called_once() + + @patch("obliteratus.capability_check.subprocess.run") + def test_raises_on_failure(self, mock_run, tmp_path): + from obliteratus.capability_check import _run_lm_eval + + mock_run.return_value = MagicMock(returncode=1, stdout="", stderr="error") + + with pytest.raises(RuntimeError, match="lm-eval exited"): + _run_lm_eval("fake/model", "mmlu", 5, "cpu", str(tmp_path)) + + +class TestCapabilityCheck: + @patch("obliteratus.capability_check._run_lm_eval") + def test_computes_delta(self, mock_lm_eval, tmp_path): + from obliteratus.capability_check import capability_check + + mock_lm_eval.side_effect = [ + {"results": {"mmlu": {"acc,none": 0.81}}}, # abliterated + {"results": {"mmlu": {"acc,none": 0.87}}}, # stock + ] + + result = capability_check( + "fake/abliterated", "fake/stock", + device="cpu", output_dir=str(tmp_path), + ) + + assert result["abliterated_acc"] == 0.81 + assert result["stock_acc"] == 0.87 + assert result["delta_pp"] == -6.0 + + @patch("obliteratus.capability_check._run_lm_eval") + def test_quick_mode(self, mock_lm_eval, tmp_path): + from obliteratus.capability_check import QUICK_SUBJECTS, capability_check + + mock_lm_eval.side_effect = [ + {"results": {s: {"acc,none": 0.8} for s in QUICK_SUBJECTS}}, + {"results": {s: {"acc,none": 0.9} for s in QUICK_SUBJECTS}}, + ] + + result = capability_check( + "fake/abliterated", "fake/stock", + device="cpu", quick=True, output_dir=str(tmp_path), + ) + + assert result["tasks"] == ",".join(QUICK_SUBJECTS) + assert abs(result["abliterated_acc"] - 0.8) < 0.01 + + @patch("obliteratus.capability_check._run_lm_eval") + def test_saves_summary(self, mock_lm_eval, tmp_path): + from obliteratus.capability_check import capability_check + + mock_lm_eval.side_effect = [ + {"results": {"mmlu": {"acc,none": 0.85}}}, + {"results": {"mmlu": {"acc,none": 0.87}}}, + ] + + capability_check( + "fake/abliterated", "fake/stock", + device="cpu", output_dir=str(tmp_path), + ) + + summary_path = tmp_path / "capability_summary.json" + assert summary_path.exists() + summary = json.loads(summary_path.read_text()) + assert "delta_pp" in summary + assert summary["method"] == "lm-eval-harness 0-shot log-likelihood"