tests/test_modality_router.py

"""Tests for VisualModalityRouter -> verifies routing rules and prompt grounding.

The router is only ever reached from an explicit student visualize request (the
old Wiki click-to-offer path was removed), so "decline" is not a valid outcome:
the router always returns one of the 5 real modalities.
"""
from unittest.mock import MagicMock

from app.agents.modality_router import VisualModalityRouter


def _make_router_with_response(modality: str) -> tuple:
    """Return (router, captured_messages, captured_kwargs) where the client returns a canned decision."""
    router = VisualModalityRouter.__new__(VisualModalityRouter)
    mock_client = MagicMock()
    router._client = mock_client

    from app.agents.modality_router import ModalityDecision
    decision = ModalityDecision(modality=modality, reasoning="test reasoning")
    captured = {}

    def fake_structured_complete(messages, model_cls, **kwargs):
        captured["messages"] = messages
        captured["kwargs"] = kwargs
        return decision

    mock_client.structured_complete.side_effect = fake_structured_complete
    return router, captured


CHUNKS_WITH_DATA = [
    {"source": "notes.pdf", "text": "Revenue in 2020: $120M, 2021: $145M, 2022: $178M"},
    {"source": "notes.pdf", "text": "Growth rate was 20.8% year over year."},
]

CHUNKS_WITHOUT_DATA = [
    {"source": "notes.pdf", "text": "Entropy is a measure of disorder in a thermodynamic system."},
    {"source": "notes.pdf", "text": "The second law states entropy increases in isolated systems."},
]


def test_classify_returns_modality_decision():
    router, _ = _make_router_with_response("graph")
    result = router.classify("revenue trends", "## Revenue\n...", CHUNKS_WITH_DATA, "high_school")
    assert result.modality == "graph"


def test_classify_uses_reasoning_effort():
    """This engine writes actual chart-shape decisions -- reasoning should be enabled."""
    router, captured = _make_router_with_response("graph")
    router.classify("revenue trends", "## Revenue\n...", CHUNKS_WITH_DATA, "high_school")
    assert captured["kwargs"].get("reasoning_effort") == "medium"


def test_prompt_contains_chunk_text():
    """The router must embed source material in the prompt so decisions are grounded."""
    router, captured = _make_router_with_response("graph")
    router.classify("revenue", "card text", CHUNKS_WITH_DATA, "graduate")
    user_message = captured["messages"][1]["content"]
    assert "Revenue in 2020" in user_message, "chunk text must appear in the user message"
    assert "notes.pdf" not in user_message  # chunk source key is not embedded, just the text


def test_prompt_contains_selection_and_level():
    router, captured = _make_router_with_response("2d_text")
    router.classify("entropy", "card", CHUNKS_WITHOUT_DATA, "eli5")
    user_message = captured["messages"][1]["content"]
    assert "entropy" in user_message
    assert "eli5" in user_message


def test_prompt_caps_chunk_text():
    """Chunks larger than 3000 chars must be truncated."""
    big_chunk = [{"source": "big.pdf", "text": "x" * 5000}]
    router, captured = _make_router_with_response("2d_text")
    router.classify("concept", "card", big_chunk, "high_school")
    user_message = captured["messages"][1]["content"]
    # The chunk text is capped at 3000 chars inside classify()
    assert len(user_message) < 5000 + 500  # well under 5000 raw chunk length


def test_2d_anim_for_dynamic_concept():
    router, _ = _make_router_with_response("2d_anim")
    result = router.classify("pendulum motion", "card", CHUNKS_WITHOUT_DATA, "high_school")
    assert result.modality == "2d_anim"


def test_graph_for_chart_request_even_with_unrelated_chunks():
    """Regression test: an explicit 'bar chart of X' request must route to graph based on
    the student's own wording, not get remapped to some fixed default because the
    retrieved chunks are unrelated to X. (Confirmed real bug: a hardcoded decline->2d_anim
    remap previously sent every request, regardless of type, to 2d_anim.)"""
    router, _ = _make_router_with_response("graph")
    result = router.classify(
        "bar chart between population of india in 1990 and in 2010",
        "card",
        CHUNKS_WITHOUT_DATA,  # deliberately unrelated to the request
        "high_school",
    )
    assert result.modality == "graph"


def test_only_five_modalities_are_valid_literal_values():
    """`decline` must not be constructible -- this router never offers it as an outcome."""
    from app.agents.modality_router import ModalityDecision
    for modality in ["formula", "graph", "2d_text", "3d", "2d_anim"]:
        decision = ModalityDecision(modality=modality, reasoning="r")
        assert decision.modality == modality

    import pytest
    from pydantic import ValidationError
    with pytest.raises(ValidationError):
        ModalityDecision(modality="decline", reasoning="r")


def test_decline_reason_field_removed():
    """decline_reason was part of the old decline outcome; it must not linger as a field."""
    from app.agents.modality_router import ModalityDecision
    assert "decline_reason" not in ModalityDecision.model_fields


def test_prompt_never_offers_decline():
    """The system prompt must explicitly rule out declining -- the student already asked
    for a visual, so the router always picks one of the 5 real shapes."""
    from app.agents.modality_router import _SYSTEM_PROMPT

    prompt_lower = _SYSTEM_PROMPT.lower()
    assert "not an option" in prompt_lower or "always pick exactly one" in prompt_lower
    for modality in ["formula", "graph", "2d_text", "3d", "2d_anim"]:
        assert modality in prompt_lower, f"modality {modality!r} missing from system prompt"