tests/test_classifier.py

"""Tests for OTSClassifier with mocked HuggingFace pipelines.

No real model downloads are needed -- the ``transformers.pipeline``
call is monkeypatched to return a controllable callable.
"""

from __future__ import annotations

from types import SimpleNamespace
from typing import Any
from unittest.mock import MagicMock

import pytest

from app.engines.semantic.classifier import (
    OTSClassifier,
    extract_positive_probability,
    resolve_positive_labels,
)
from app.engines.semantic.config import ModelConfig
from app.engines.semantic.models import SemanticEngineId
from app.engines.semantic.truncation import HeadOnly


# ---- Fixtures ---------------------------------------------------------


@pytest.fixture
def simple_config() -> ModelConfig:
    """Minimal ModelConfig for testing."""
    return ModelConfig(
        model_id="test/model",
        positive_label_ids=(1,),
        max_tokens=480,
        max_length=512,
        head_ratio=0.67,
    )


@pytest.fixture
def multi_class_config() -> ModelConfig:
    """ModelConfig where multiple label indices are positive."""
    return ModelConfig(
        model_id="test/multi-model",
        positive_label_ids=(1, 3),
        max_tokens=480,
        max_length=512,
        head_ratio=0.67,
    )


def _make_mock_pipeline(
    output: list[dict[str, Any]],
    id2label: dict[int, str] | None = None,
) -> MagicMock:
    """Build a mock that quacks like ``transformers.pipeline(...)``."""
    pipe = MagicMock()
    pipe.return_value = output
    # Model config for label resolution
    config = SimpleNamespace(id2label=id2label or {})
    pipe.model = SimpleNamespace(config=config)
    return pipe


# ---- resolve_positive_labels -----------------------------------------


class TestResolvePositiveLabels:
    """Label name resolution from integer indices."""

    def test_generic_label_format(self):
        """Always includes LABEL_N and str(N) variants."""
        pipe = _make_mock_pipeline([], id2label={})
        labels = resolve_positive_labels(pipe, (1,))
        assert "LABEL_1" in labels
        assert "1" in labels

    def test_id2label_mapping(self):
        """Includes model-specific label names from id2label."""
        pipe = _make_mock_pipeline([], id2label={1: "INJECTION"})
        labels = resolve_positive_labels(pipe, (1,))
        assert "INJECTION" in labels
        assert "LABEL_1" in labels

    def test_multiple_positive_ids(self):
        """All indices contribute label variants."""
        pipe = _make_mock_pipeline(
            [], id2label={1: "phishing", 3: "spam"},
        )
        labels = resolve_positive_labels(pipe, (1, 3))
        assert "LABEL_1" in labels
        assert "LABEL_3" in labels
        assert "phishing" in labels
        assert "spam" in labels

    def test_no_model_config(self):
        """Handles pipelines without model.config gracefully."""
        pipe = MagicMock()
        pipe.model = None
        labels = resolve_positive_labels(pipe, (1,))
        assert "LABEL_1" in labels
        assert "1" in labels


# ---- extract_positive_probability ------------------------------------


class TestExtractPositiveProbability:
    """Probability extraction from pipeline output."""

    def test_flat_list(self):
        """Standard flat list output."""
        output = [
            {"label": "LABEL_0", "score": 0.3},
            {"label": "LABEL_1", "score": 0.7},
        ]
        prob = extract_positive_probability(output, {"LABEL_1"})
        assert prob == pytest.approx(0.7)

    def test_nested_list(self):
        """Doubly-nested list output (some HF pipeline versions)."""
        output = [[
            {"label": "LABEL_0", "score": 0.2},
            {"label": "LABEL_1", "score": 0.8},
        ]]
        prob = extract_positive_probability(output, {"LABEL_1"})
        assert prob == pytest.approx(0.8)

    def test_multi_class_sum(self):
        """Probabilities of multiple positive classes are summed."""
        output = [
            {"label": "LABEL_0", "score": 0.1},
            {"label": "LABEL_1", "score": 0.4},
            {"label": "LABEL_2", "score": 0.2},
            {"label": "LABEL_3", "score": 0.3},
        ]
        prob = extract_positive_probability(output, {"LABEL_1", "LABEL_3"})
        assert prob == pytest.approx(0.7)

    def test_no_positive_labels_found_raises(self):
        """Raises ValueError when no output labels match the positive set."""
        output = [{"label": "LABEL_0", "score": 1.0}]
        with pytest.raises(ValueError, match="No output labels matched"):
            extract_positive_probability(output, {"LABEL_1"})

    def test_clamped_to_one(self):
        """Result is clamped to 1.0 even if scores sum higher."""
        output = [
            {"label": "LABEL_1", "score": 0.6},
            {"label": "LABEL_3", "score": 0.6},
        ]
        prob = extract_positive_probability(output, {"LABEL_1", "LABEL_3"})
        assert prob == 1.0

    def test_model_specific_labels(self):
        """Matches model-specific string labels (e.g. 'phishing')."""
        output = [
            {"label": "benign", "score": 0.2},
            {"label": "phishing", "score": 0.8},
        ]
        prob = extract_positive_probability(output, {"phishing"})
        assert prob == pytest.approx(0.8)


# ---- _validate_positive_label_ids ------------------------------------


class TestValidatePositiveLabelIds:
    """Fail-fast validation during load()."""

    def test_valid_ids_pass(self):
        """No error when all positive IDs exist in id2label."""
        from app.engines.semantic.classifier import _validate_positive_label_ids

        pipe = _make_mock_pipeline([], id2label={0: "benign", 1: "INJECTION"})
        # Should not raise
        _validate_positive_label_ids(pipe, (1,))

    def test_missing_id_raises(self):
        """ValueError when a positive ID is absent from id2label."""
        from app.engines.semantic.classifier import _validate_positive_label_ids

        pipe = _make_mock_pipeline([], id2label={0: "benign", 1: "INJECTION"})
        with pytest.raises(ValueError, match="positive_label_ids contains 3"):
            _validate_positive_label_ids(pipe, (1, 3))

    def test_empty_id2label_skips_validation(self):
        """No validation when id2label is empty (LABEL_N fallback used)."""
        from app.engines.semantic.classifier import _validate_positive_label_ids

        pipe = _make_mock_pipeline([], id2label={})
        # Should not raise -- no id2label to validate against
        _validate_positive_label_ids(pipe, (1, 3))

    def test_no_model_config_skips_validation(self):
        """No validation when model has no config attribute."""
        from app.engines.semantic.classifier import _validate_positive_label_ids

        pipe = MagicMock()
        pipe.model = None
        # Should not raise
        _validate_positive_label_ids(pipe, (1,))


# ---- OTSClassifier ---------------------------------------------------


class TestOTSClassifier:
    """OTSClassifier lifecycle and prediction tests."""

    def test_predict_before_load_raises(self, simple_config):
        """Calling predict() before load() raises RuntimeError."""
        clf = OTSClassifier(
            engine_id=SemanticEngineId.PROMPT_INJECTION,
            config=simple_config,
            threshold=0.5,
        )
        with pytest.raises(RuntimeError, match="not loaded"):
            clf.predict("test text")

    def test_predict_returns_prediction(self, monkeypatch, simple_config):
        """After load(), predict() returns a valid SemanticPrediction."""
        mock_pipe = _make_mock_pipeline(
            [{"label": "LABEL_0", "score": 0.3}, {"label": "LABEL_1", "score": 0.7}],
            id2label={0: "benign", 1: "INJECTION"},
        )

        import sys

        fake_transformers = SimpleNamespace(
            pipeline=lambda *a, **kw: mock_pipe,
        )
        monkeypatch.setitem(sys.modules, "transformers", fake_transformers)

        clf = OTSClassifier(
            engine_id=SemanticEngineId.PROMPT_INJECTION,
            config=simple_config,
            threshold=0.5,
        )
        clf.load()
        pred = clf.predict("test text")

        assert pred.engine_id == SemanticEngineId.PROMPT_INJECTION
        assert pred.probability == pytest.approx(0.7)
        assert pred.predicted_label == 1
        assert pred.threshold == 0.5
        assert pred.model_id == "test/model"

    def test_threshold_determines_label(self, monkeypatch, simple_config):
        """Label is 0 when probability < threshold."""
        mock_pipe = _make_mock_pipeline(
            [{"label": "LABEL_0", "score": 0.7}, {"label": "LABEL_1", "score": 0.3}],
        )

        import sys

        fake_transformers = SimpleNamespace(
            pipeline=lambda *a, **kw: mock_pipe,
        )
        monkeypatch.setitem(sys.modules, "transformers", fake_transformers)

        clf = OTSClassifier(
            engine_id=SemanticEngineId.PROMPT_INJECTION,
            config=simple_config,
            threshold=0.5,
        )
        clf.load()
        pred = clf.predict("test text")

        assert pred.probability == pytest.approx(0.3)
        assert pred.predicted_label == 0

    def test_truncation_applied(self, monkeypatch, simple_config):
        """Text is truncated before being sent to the pipeline."""
        called_with: list[str] = []

        def capturing_pipe(text, **kwargs):
            called_with.append(text)
            return [
                {"label": "LABEL_0", "score": 0.5},
                {"label": "LABEL_1", "score": 0.5},
            ]

        mock_pipe = MagicMock(side_effect=capturing_pipe)
        mock_pipe.model = SimpleNamespace(config=SimpleNamespace(id2label={}))

        import sys

        fake_transformers = SimpleNamespace(
            pipeline=lambda *a, **kw: mock_pipe,
        )
        monkeypatch.setitem(sys.modules, "transformers", fake_transformers)

        clf = OTSClassifier(
            engine_id=SemanticEngineId.PROMPT_INJECTION,
            config=simple_config,
            threshold=0.5,
            truncation=HeadOnly(max_tokens=3),
        )
        clf.load()
        clf.predict("one two three four five")

        assert len(called_with) == 1
        assert called_with[0] == "one two three"

    def test_multi_class_positive_sum(self, monkeypatch, multi_class_config):
        """Multiple positive label IDs have their probabilities summed."""
        mock_pipe = _make_mock_pipeline(
            [
                {"label": "LABEL_0", "score": 0.1},
                {"label": "LABEL_1", "score": 0.3},
                {"label": "LABEL_2", "score": 0.2},
                {"label": "LABEL_3", "score": 0.4},
            ],
        )

        import sys

        fake_transformers = SimpleNamespace(
            pipeline=lambda *a, **kw: mock_pipe,
        )
        monkeypatch.setitem(sys.modules, "transformers", fake_transformers)

        clf = OTSClassifier(
            engine_id=SemanticEngineId.MALICIOUS_INTENT,
            config=multi_class_config,
            threshold=0.5,
        )
        clf.load()
        pred = clf.predict("test")

        # Labels 1 + 3: 0.3 + 0.4 = 0.7
        assert pred.probability == pytest.approx(0.7)
        assert pred.predicted_label == 1

    def test_revision_passed_to_pipeline(self, monkeypatch, simple_config):
        """The config's revision is forwarded to the HF pipeline call."""
        captured_kwargs: list[dict] = []

        def capturing_factory(*args, **kwargs):
            captured_kwargs.append(kwargs)
            return _make_mock_pipeline(
                [{"label": "LABEL_0", "score": 0.5}, {"label": "LABEL_1", "score": 0.5}],
            )

        import sys

        fake_transformers = SimpleNamespace(pipeline=capturing_factory)
        monkeypatch.setitem(sys.modules, "transformers", fake_transformers)

        config_with_rev = ModelConfig(
            model_id="test/model",
            positive_label_ids=(1,),
            max_tokens=480,
            max_length=512,
            head_ratio=0.67,
            revision="abc123",
        )
        clf = OTSClassifier(
            engine_id=SemanticEngineId.PROMPT_INJECTION,
            config=config_with_rev,
            threshold=0.5,
        )
        clf.load()

        assert len(captured_kwargs) == 1
        assert captured_kwargs[0]["revision"] == "abc123"

    def test_load_validates_positive_label_ids(self, monkeypatch):
        """load() raises ValueError if positive IDs are absent from id2label."""
        mock_pipe = _make_mock_pipeline(
            [],
            id2label={0: "safe", 1: "unsafe"},
        )

        import sys

        fake_transformers = SimpleNamespace(
            pipeline=lambda *a, **kw: mock_pipe,
        )
        monkeypatch.setitem(sys.modules, "transformers", fake_transformers)

        bad_config = ModelConfig(
            model_id="test/bad",
            positive_label_ids=(1, 5),  # 5 doesn't exist
            max_tokens=480,
            max_length=512,
            head_ratio=0.67,
        )
        clf = OTSClassifier(
            engine_id=SemanticEngineId.PROMPT_INJECTION,
            config=bad_config,
            threshold=0.5,
        )
        with pytest.raises(ValueError, match="positive_label_ids contains 5"):
            clf.load()