OBLITERATUS/tests/test_cli.py

"""CLI dispatch tests for obliteratus.cli.main().

These tests verify argument parsing and subcommand routing without
downloading real models or running any pipeline.  They use
``unittest.mock.patch`` to capture stdout/stderr and
``pytest.raises(SystemExit)`` for argparse exits.
"""

from __future__ import annotations

import json
from io import StringIO
from types import SimpleNamespace
from unittest.mock import MagicMock
from unittest.mock import patch

import pytest

from obliteratus.cli import main


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def _capture_exit(argv: list[str] | None, *, expect_code: int | None = None):
    """Call main(argv), expecting SystemExit; return captured stderr text."""
    buf = StringIO()
    with pytest.raises(SystemExit) as exc_info, patch("sys.stderr", buf):
        main(argv)
    if expect_code is not None:
        assert exc_info.value.code == expect_code
    return buf.getvalue()


# ---------------------------------------------------------------------------
# Tests
# ---------------------------------------------------------------------------


class TestCLIDispatch:
    """Test suite for CLI argument parsing and subcommand dispatch."""

    # 1. No args -> prints help / exits with error
    def test_main_no_args_prints_help(self):
        """Calling main() with no args should exit (subcommand is required)."""
        stderr_text = _capture_exit([], expect_code=2)
        # argparse prints usage info to stderr on error
        assert "usage" in stderr_text.lower() or "required" in stderr_text.lower()

    # 2. models command lists models without error
    def test_models_command(self):
        """Calling main(['models']) should list models without raising."""
        with patch("obliteratus.cli.console") as mock_console:
            main(["models"])
        # console.print is called at least once to render the table
        assert mock_console.print.call_count >= 1

    # 3. obliterate without model arg -> error
    def test_obliterate_requires_model(self):
        """Calling main(['obliterate']) without a model arg should error."""
        stderr_text = _capture_exit(["obliterate"], expect_code=2)
        assert "model" in stderr_text.lower() or "required" in stderr_text.lower()

    # 4. obliterate --method accepts valid methods
    def test_obliterate_valid_methods(self):
        """Test that --method accepts all 9 pipeline methods."""
        valid_methods = [
            "basic", "advanced", "aggressive", "spectral_cascade",
            "informed", "surgical", "optimized", "inverted", "nuclear",
        ]
        for method in valid_methods:
            # Patch the actual pipeline execution so nothing runs
            with patch("obliteratus.cli._cmd_abliterate") as mock_cmd:
                main(["obliterate", "fake/model", "--method", method])
                mock_cmd.assert_called_once()
                args_passed = mock_cmd.call_args[0][0]
                assert args_passed.method == method

    # 4b. invalid methods are rejected
    def test_obliterate_rejects_invalid_method(self):
        """The CLI --method flag rejects unknown method names."""
        stderr_text = _capture_exit(
            ["obliterate", "fake/model", "--method", "nonexistent"],
            expect_code=2,
        )
        assert "invalid choice" in stderr_text.lower()

    # 5. run requires config path
    def test_run_requires_config(self):
        """Calling main(['run']) without a config path should error."""
        stderr_text = _capture_exit(["run"], expect_code=2)
        assert "config" in stderr_text.lower() or "required" in stderr_text.lower()

    # 6. aggregate with nonexistent dir handles gracefully
    def test_aggregate_command_missing_dir(self):
        """Calling main(['aggregate']) with nonexistent dir should handle gracefully."""
        with patch("obliteratus.cli.console") as mock_console:
            main(["aggregate", "--dir", "/nonexistent/path/to/nowhere"])
        # The command prints a message about no contributions found and returns
        printed_text = " ".join(
            str(call) for call in mock_console.print.call_args_list
        )
        assert "no contributions found" in printed_text.lower() or mock_console.print.called

    def test_aggregate_accepts_format_metric_min_runs(self):
        """aggregate accepts --format, --metric and --min-runs flags."""
        with patch("obliteratus.cli._cmd_aggregate") as mock_cmd:
            main([
                "aggregate",
                "--format", "latex",
                "--metric", "refusal_rate",
                "--min-runs", "3",
            ])
            mock_cmd.assert_called_once()
            args_passed = mock_cmd.call_args[0][0]
            assert args_passed.format == "latex"
            assert args_passed.metric == "refusal_rate"
            assert args_passed.min_runs == 3

    def test_aggregate_rejects_invalid_format(self):
        """aggregate rejects unknown --format choices."""
        stderr_text = _capture_exit(
            ["aggregate", "--format", "invalid"],
            expect_code=2,
        )
        assert "invalid choice" in stderr_text.lower()

    # 7. --help flag prints help
    def test_help_flag(self):
        """Calling main(['--help']) should print help and exit 0."""
        buf = StringIO()
        with pytest.raises(SystemExit) as exc_info, patch("sys.stdout", buf):
            main(["--help"])
        assert exc_info.value.code == 0
        output = buf.getvalue()
        assert "obliteratus" in output.lower() or "usage" in output.lower()

    # 8. interactive subcommand is registered
    def test_interactive_command_exists(self):
        """Verify 'interactive' subcommand is registered and dispatches."""
        with patch("obliteratus.cli._cmd_interactive") as mock_cmd:
            main(["interactive"])
            mock_cmd.assert_called_once()

    # 9. --contribute and --contribute-notes are accepted on obliterate
    def test_contribute_flags_on_obliterate(self):
        """Verify --contribute and --contribute-notes are accepted args."""
        with patch("obliteratus.cli._cmd_abliterate") as mock_cmd:
            main([
                "obliterate", "fake/model",
                "--contribute",
                "--contribute-notes", "Testing contribution system",
            ])
            mock_cmd.assert_called_once()
            args_passed = mock_cmd.call_args[0][0]
            assert args_passed.contribute is True
            assert args_passed.contribute_notes == "Testing contribution system"

    @pytest.mark.parametrize("flag", ["--prompt-pairs-file", "--prompt-pair-file"])
    @pytest.mark.parametrize("command", ["obliterate", "abliterate"])
    def test_prompt_pairs_file_flag_is_available_on_obliterate_and_alias(
        self,
        command,
        flag,
        tmp_path,
    ):
        """Explicit prompt-pair files are parsed for both commands and flag spellings."""
        path = tmp_path / "pairs.json"
        path.write_text("{}", encoding="utf-8")
        with patch("obliteratus.cli._cmd_abliterate") as mock_cmd:
            main([command, "fake/model", flag, str(path)])
            args_passed = mock_cmd.call_args[0][0]
            assert args_passed.prompt_pairs_file == str(path)

    @pytest.mark.parametrize("flag", ["--prompt-pairs-file", "--prompt-pair-file"])
    @pytest.mark.parametrize("command", ["obliterate", "abliterate"])
    def test_prompt_pairs_file_is_mutually_exclusive_with_residue_files(
        self,
        command,
        flag,
        tmp_path,
    ):
        """Explicit prompt-pair files and mined residue construction cannot be mixed."""
        path = tmp_path / "pairs.json"
        path.write_text("{}", encoding="utf-8")
        stderr_text = _capture_exit(
            [
                command,
                "fake/model",
                flag,
                str(path),
                "--residue-file",
                "audit.json",
            ],
            expect_code=2,
        )
        assert "not allowed with argument" in stderr_text.lower()

    @pytest.mark.parametrize(
        ("option", "value"),
        [
            ("--dataset", "custom"),
            ("--residue-weight", "7"),
            ("--residue-max", "3"),
        ],
    )
    @pytest.mark.parametrize("flag", ["--prompt-pairs-file", "--prompt-pair-file"])
    @pytest.mark.parametrize("command", ["obliterate", "abliterate"])
    def test_prompt_pairs_file_rejects_residue_only_options(
        self,
        command,
        flag,
        option,
        value,
        tmp_path,
    ):
        """Explicit prompt-pair files cannot be mixed with residue-only options."""
        path = tmp_path / "pairs.json"
        path.write_text("{}", encoding="utf-8")
        stderr_text = _capture_exit(
            [command, "fake/model", flag, str(path), option, value],
            expect_code=2,
        )
        assert option in stderr_text
        assert "can only be used with --residue-file" in stderr_text


def test_cmd_abliterate_wires_prompt_pairs_file_into_pipeline(tmp_path):
    """Loaded prompt-pair files are passed directly to AbliterationPipeline."""
    path = tmp_path / "pairs.json"
    path.write_text(
        json.dumps(
            {
                "harmful": [f"harm {index}" for index in range(5)],
                "harmless": [f"safe {index}" for index in range(5)],
            }
        ),
        encoding="utf-8",
    )
    result_path = tmp_path / "result"
    result_path.mkdir()
    pipeline = MagicMock()
    pipeline.run.return_value = str(result_path)

    class FakeLive:
        def __init__(self, *_args, **_kwargs):
            self.update = MagicMock()

        def __enter__(self):
            return self

        def __exit__(self, *_args):
            return False

    args = SimpleNamespace(
        model="org/model",
        output_dir=str(tmp_path / "out"),
        device="cpu",
        dtype="float32",
        method="basic",
        n_directions=1,
        direction_method=None,
        regularization=None,
        refinement_passes=1,
        min_layer_fraction=None,
        max_layer_fraction=None,
        harmless_pc_count=None,
        shield_concept_count=None,
        shield_ridge=None,
        shield_residualize=None,
        shield_layer_penalty=None,
        projection_target=None,
        projection_row_fraction=None,
        quantization=None,
        gpu_memory_utilization=None,
        large_model=False,
        verify_sample_size=1,
        refusal_max_tokens=1,
        residue_file=[],
        dataset="builtin",
        residue_weight=5,
        residue_max=None,
        prompt_pairs_file=str(path),
        contribute=False,
        contribute_notes="",
    )
    with (
        patch("rich.live.Live", FakeLive),
        patch("obliteratus.abliterate.AbliterationPipeline", return_value=pipeline) as factory,
        patch("obliteratus.telemetry.maybe_send_pipeline_report"),
    ):
        from obliteratus import cli

        cli._cmd_abliterate(args)

    assert factory.call_args.kwargs["harmful_prompts"] == [f"harm {index}" for index in range(5)]
    assert factory.call_args.kwargs["harmless_prompts"] == [f"safe {index}" for index in range(5)]

    args.prompt_pairs_file = None
    pipeline.reset_mock()
    pipeline.run.return_value = str(result_path)
    with (
        patch("rich.live.Live", FakeLive),
        patch("obliteratus.abliterate.AbliterationPipeline", return_value=pipeline) as factory,
        patch("obliteratus.telemetry.maybe_send_pipeline_report"),
    ):
        from obliteratus import cli

        cli._cmd_abliterate(args)

    assert "harmful_prompts" not in factory.call_args.kwargs
    assert "harmless_prompts" not in factory.call_args.kwargs


def test_cmd_abliterate_reports_invalid_prompt_pairs_file(tmp_path):
    from obliteratus import cli

    args = SimpleNamespace(
        model="org/model",
        output_dir=str(tmp_path / "out"),
        device="cpu",
        dtype="float32",
        method="basic",
        n_directions=1,
        direction_method=None,
        regularization=None,
        refinement_passes=1,
        min_layer_fraction=None,
        max_layer_fraction=None,
        harmless_pc_count=None,
        shield_concept_count=None,
        shield_ridge=None,
        shield_residualize=None,
        shield_layer_penalty=None,
        projection_target=None,
        projection_row_fraction=None,
        quantization=None,
        gpu_memory_utilization=None,
        large_model=False,
        verify_sample_size=1,
        refusal_max_tokens=1,
        residue_file=[],
        dataset="builtin",
        residue_weight=5,
        residue_max=None,
        prompt_pairs_file=str(tmp_path / "missing.json"),
        contribute=False,
        contribute_notes="",
    )

    with pytest.raises(SystemExit) as exc:
        cli._cmd_abliterate(args)

    assert exc.value.code == 2


class _EncodingOnlyStdout:
    """Minimal stream stand-in for encoding-selection tests."""

    def __init__(self, encoding: str | None) -> None:
        self.encoding = encoding


class TestConsoleEncoding:
    """The CLI must remain renderable on legacy Windows code pages."""

    @pytest.mark.parametrize("encoding", ["cp1252", "ascii", "not-a-codec"])
    def test_console_text_falls_back_when_text_is_not_encodable(self, encoding):
        from obliteratus.cli import _console_text

        with patch("sys.stdout", _EncodingOnlyStdout(encoding)):
            assert _console_text("█→", "FALLBACK") == "FALLBACK"

    @pytest.mark.parametrize("encoding", ["utf-8", None])
    def test_console_text_keeps_unicode_on_utf8_compatible_streams(self, encoding):
        from obliteratus.cli import _console_text

        with patch("sys.stdout", _EncodingOnlyStdout(encoding)):
            assert _console_text("█→", "FALLBACK") == "█→"

    def test_banner_degrades_to_ascii_on_cp1252(self):
        from obliteratus.cli import _banner_for_console

        with patch("sys.stdout", _EncodingOnlyStdout("cp1252")):
            banner = _banner_for_console()

        assert "OBLITERATUS" in banner
        assert banner.isascii()

    def test_main_renders_selected_banner(self):
        with (
            patch("sys.stdout", _EncodingOnlyStdout("cp1252")),
            patch("obliteratus.cli.console") as mock_console,
        ):
            main(["models"])

        assert mock_console.print.call_args_list[0].args[0].isascii()