tests/agents/d3/test_router.py

from app.agents.d3.router import D3TemplateRouter, D3Selection


class _FakeClient:
    def __init__(self, obj): self._obj = obj
    def structured_complete(self, messages, schema, **kw): return self._obj


def test_router_picks_template():
    r = D3TemplateRouter(client=_FakeClient(D3Selection(template_id="bar_chart")))
    out = r.select("GDP by country", [{"source": "x", "text": "US 21, China 14"}], "graduate")
    assert out.template_id == "bar_chart"


def test_router_never_declines_falls_back_on_unknown_template_id():
    # The router never declines -- the student already asked for a chart. If the
    # model returns a template id that doesn't exist in the catalog, fall back to
    # a safe default instead of refusing.
    r = D3TemplateRouter(client=_FakeClient(D3Selection(template_id="not_a_real_template")))
    out = r.select("the nature of justice", [{"source": "x", "text": "philosophy"}], "graduate")
    assert out.template_id == "bar_chart"