stack/tests/llm/test_api.py

"""llm.api — FastAPI chat app."""

from unittest.mock import patch

from fastapi.testclient import TestClient

from llm.api import app

client = TestClient(app)


class TestHealth:
    def test_ok(self):
        r = client.get("/health")
        assert r.status_code == 200
        assert r.json() == {"status": "ok"}


class TestIndex:
    def test_serves_chat_page(self):
        r = client.get("/")
        assert r.status_code == 200
        assert "text/html" in r.headers["content-type"]
        assert "Library Chat" in r.text
        assert 'id="form"' in r.text

    def test_serves_chat_page_with_valuation_renderer(self):
        r = client.get("/")
        assert r.status_code == 200
        html = r.text
        assert "function renderValuation" in html
        assert "ev.type === 'valuation'" in html
        assert "table.valuation" in html
        # the money on screen is national and unadjusted — say so
        assert "no geographic (GPCI) adjustment" in html
        assert "createCaption" in html
        # notes explain an unpaid status and which CF is shown
        assert "r.status_note" in html and "r.cf_note" in html
        assert "'cf', 'CF (non-APM)'" in html
        # a citation without a URL is not an empty <a href="">
        assert "if (p.url)" in html

    def test_serves_chat_page_with_lineage_renderer(self):
        r = client.get("/")
        assert r.status_code == 200
        html = r.text
        assert "function renderLineage" in html
        assert "ev.type === 'lineage'" in html
        assert "table.lineage" in html
        assert 'id="mode"' in html
        assert "tr.unanchored" in html
        assert "ul.className = 'elements'" in html
        assert "ul.className = 'guidance'" in html
        # mode is sent to the API and echoed back on the sources meta line
        assert "mode" in html and "ev.mode === 'timeline'" in html

    def test_serves_chat_page_with_markdown_and_cite_links(self):
        """The answer is painted as markdown (headings/lists/bold, never
        raw "## …"), and a cited [label] whose URL arrived on the
        lineage/valuation/sources events is an <a class="cite"> jump
        link to the FR paragraph / eCFR section; the answer is repainted
        when the sources event completes the label→url map."""
        html = client.get("/").text
        assert "function paintInline" in html
        assert "function paint(el, text, links)" in html
        assert "function collectLinks" in html
        assert "a.cite" in html
        assert (
            "collectLinks(links, ev); renderSources(wrap, ev); paint(b, answer, links)"
            in html
        )
        assert "document.createElement('h' + level)" in html
        # model text never reaches innerHTML
        assert "innerHTML = ''" in html and "innerHTML = text" not in html


class TestWhoami:
    def test_reads_forwarded_header(self):
        r = client.get("/whoami", headers={"X-Auth-Request-User": "kert"})
        assert r.json() == {"user": "kert"}

    def test_empty_when_absent(self):
        r = client.get("/whoami")
        assert r.json() == {"user": ""}


class TestChat:
    def test_empty_question_400(self):
        r = client.post("/chat", json={"question": "   "})
        assert r.status_code == 400

    @patch("llm.rag.stream_answer")
    def test_streams_sse_events(self, mock_stream):
        mock_stream.return_value = iter(
            [
                {"type": "token", "text": "Hi"},
                {"type": "sources", "sources": []},
                {"type": "done"},
            ]
        )
        r = client.post("/chat", json={"question": "hello"})
        assert r.status_code == 200
        assert "text/event-stream" in r.headers["content-type"]
        body = r.text
        assert 'data: {"type": "token", "text": "Hi"}' in body
        assert '"type": "done"' in body

    @patch("llm.rag.stream_answer")
    def test_error_becomes_event_not_500(self, mock_stream):
        mock_stream.side_effect = RuntimeError("ollama down")
        r = client.post("/chat", json={"question": "hello"})
        assert r.status_code == 200
        assert '"type": "error"' in r.text
        assert "ollama down" in r.text


class TestChatSince:
    @patch("llm.rag.stream_answer")
    def test_since_forwarded(self, mock_stream):
        mock_stream.return_value = iter([{"type": "done"}])
        r = client.post("/chat", json={"question": "q", "since": "2025-09-01"})
        assert r.status_code == 200
        assert mock_stream.call_args.kwargs["since"] == "2025-09-01"

    def test_bad_since_400(self):
        r = client.post("/chat", json={"question": "q", "since": "last year"})
        assert r.status_code == 400


class TestChatMode:
    @patch("llm.rag.stream_answer")
    def test_mode_forwarded(self, mock_stream):
        mock_stream.return_value = iter([{"type": "done"}])
        r = client.post("/chat", json={"question": "q", "mode": "timeline"})
        assert r.status_code == 200
        assert mock_stream.call_args.kwargs["mode"] == "timeline"

    @patch("llm.rag.stream_answer")
    def test_mode_defaults_to_auto(self, mock_stream):
        mock_stream.return_value = iter([{"type": "done"}])
        r = client.post("/chat", json={"question": "q"})
        assert r.status_code == 200
        assert mock_stream.call_args.kwargs["mode"] == "auto"

    def test_bad_mode_400(self):
        r = client.post("/chat", json={"question": "q", "mode": "x"})
        assert r.status_code == 400


def _cfg(**kw):
    from llm.config import LlmConfig

    base = dict(
        ollama_hosts=("http://h1:11434", "http://h2:11434"),
        host_vram={"http://h1:11434": 12, "http://h2:11434": 24},
        embed_model="e",
        instruct_model="chat",
        embed_dim=768,
        build_ann_index=False,
        pg_host="x",
        pg_port=5432,
        pg_db="llm",
        pg_user="llm",
    )
    base.update(kw)
    return LlmConfig(**base)


class TestStartup:
    """#699 ruling B3: the replica is warmed at boot, not on the chat's
    first valuation lookup. ``with TestClient(app) as c:`` is required to
    actually trigger FastAPI's startup event — a bare ``TestClient(app)``
    (as ``client`` above, module-level) never sends the ASGI lifespan
    ``startup`` message."""

    @patch("llm.config.load")
    def test_startup_event_warms_the_replica(self, mock_load):
        mock_load.return_value = _cfg()
        with patch("llm.evidence.warm") as mock_warm:
            with TestClient(app):
                pass
        mock_warm.assert_called_once_with(mock_load.return_value)


class TestHosts:
    @patch("llm.pool.pick_model", return_value="big")
    @patch("llm.pool.HostPool.check", return_value=["http://h2:11434"])
    @patch("llm.pool.HostPool.status")
    @patch("llm.config.load")
    def test_reports_fleet_and_pick(self, mock_load, mock_status, _check, _pm):
        mock_load.return_value = _cfg()
        mock_status.return_value = [
            {"host": "http://h2:11434", "vram_gb": 24.0, "models": ["chat:latest"]}
        ]
        r = client.get("/hosts")
        assert r.status_code == 200
        body = r.json()
        assert body["generation"] == {"host": "http://h2:11434", "model": "big"}
        assert [(h["host"], h["live"]) for h in body["hosts"]] == [
            ("http://h1:11434", False),
            ("http://h2:11434", True),
        ]

    @patch("llm.pool.HostPool.check", side_effect=RuntimeError("no Ollama host"))
    @patch("llm.config.load")
    def test_no_live_hosts(self, mock_load, _check):
        mock_load.return_value = _cfg(ollama_hosts=("http://h1:11434",))
        r = client.get("/hosts")
        assert r.json()["generation"] is None
        assert "no Ollama host" in r.json()["error"]