tests/core/test_route.py

import pytest

from core.ports import ChatModel
from core.schemas import RouteResponse
from core.route import (
    route_from_chips,
    route_from_message,
    resolve_route,
    step1_trace_line,
    Step1Result,
)


class StubChatModel(ChatModel):
    """Records the schema it was called with; returns a canned RouteResponse."""

    def __init__(self, canned: RouteResponse):
        self._canned = canned
        self.calls = []

    def structured(self, messages, schema):
        self.calls.append((messages, schema))
        return self._canned

    def complete(self, messages):  # pragma: no cover - unused here
        raise NotImplementedError


def test_route_from_chips_is_identity():
    r = route_from_chips(["WASH"], ["Humanitarian Conditions"], ["Context"])
    assert isinstance(r, RouteResponse)
    assert r.sectors == ["WASH"]
    assert r.lens == ["Humanitarian Conditions"]
    assert r.type == ["Context"]


def test_route_from_message_calls_model_with_route_schema():
    canned = RouteResponse(sectors=["WASH"], lens=[], type=[])
    stub = StubChatModel(canned)
    r = route_from_message("water access in Aleppo?", stub)
    assert r is canned
    assert len(stub.calls) == 1
    _messages, schema = stub.calls[0]
    assert schema is RouteResponse


def test_route_from_message_prompt_contains_axis_labels():
    stub = StubChatModel(RouteResponse(sectors=[], lens=[], type=[]))
    route_from_message("q", stub)
    messages, _schema = stub.calls[0]
    blob = " ".join(m["content"] for m in messages)
    assert "WASH" in blob and "Humanitarian Conditions" in blob and "Context" in blob


def test_resolve_route_message_partitions_out_of_scope():
    canned = RouteResponse(sectors=["WASH", "Health"], lens=["Impact"], type=[])
    res = resolve_route("message", "q", [], [], [], StubChatModel(canned))
    assert isinstance(res, Step1Result)
    assert res.route.sectors == ["WASH", "Health"]   # full faithful list preserved
    assert res.in_scope == ["WASH"]
    assert res.out_of_scope == ["Health"]


def test_resolve_route_chips_no_model_needed():
    res = resolve_route("chips", "", ["Food Security"], ["At Risk"], ["Context"], None)
    assert res.route.sectors == ["Food Security"]
    assert res.in_scope == ["Food Security"]
    assert res.out_of_scope == []


def test_resolve_route_empty_question_raises():
    with pytest.raises(ValueError, match="question"):
        resolve_route("message", "   ", [], [], [], StubChatModel(
            RouteResponse(sectors=[], lens=[], type=[])))


def test_resolve_route_zero_in_scope_passes_through():
    canned = RouteResponse(sectors=["Health"], lens=[], type=[])
    res = resolve_route("message", "q", [], [], [], StubChatModel(canned))
    assert res.in_scope == []
    assert res.out_of_scope == ["Health"]


def test_resolve_route_chips_zero_in_scope_passes_through():
    res = resolve_route("chips", "", [], ["At Risk"], [], None)
    assert res.in_scope == []


def test_step1_trace_line_with_lens():
    res = Step1Result(
        route=RouteResponse(sectors=["WASH", "Food Security"], lens=["Humanitarian Conditions"], type=[]),
        in_scope=["WASH", "Food Security"],
        out_of_scope=[],
    )
    line = step1_trace_line(res)
    assert line == "Step 1 — Understood request → WASH, Food Security (lens: Humanitarian Conditions)"


def test_step1_trace_line_without_lens():
    res = Step1Result(
        route=RouteResponse(sectors=["WASH"], lens=[], type=[]),
        in_scope=["WASH"],
        out_of_scope=["Health"],
    )
    assert step1_trace_line(res) == "Step 1 — Understood request → WASH, Health"