stack/tests/llm/test_index_refs.py

from __future__ import annotations

from unittest.mock import MagicMock, patch

from llm.chunk import Doc, content_hash
from llm.config import LlmConfig
from llm.index import index_refs
from llm.source import DocRef

CFG = LlmConfig(
    ollama_hosts=("http://h1:11434",),
    embed_model="m",
    instruct_model="g",
    embed_dim=768,
    build_ann_index=False,
    pg_host="x",
    pg_port=5432,
    pg_db="llm",
    pg_user="llm",
)
DOC = Doc(key="K1", text="Some body text.", metadata={"docket": "D"})


def _ref(fp="fp1", docket="D", loads=None, doc=DOC):
    def _load():
        if loads is not None:
            loads.append(doc.key)
        return doc

    return DocRef(
        key=doc.key, collection="comments", docket=docket, fingerprint=fp, load=_load
    )


def _run(
    refs,
    state_rows,
    *,
    force=False,
    sealed=None,
    mark_complete=True,
    complete_rows=(),
    chunks_for=lambda d, c: c,
):
    store = MagicMock()
    engine = MagicMock()
    conn = engine.begin.return_value.__enter__.return_value

    def fake_execute(clause, *a, **k):
        sql = str(clause)
        r = MagicMock()
        if "FROM index_docket_state" in sql:
            r.fetchall.return_value = list(complete_rows)
        elif "FROM index_state" in sql:
            r.fetchall.return_value = state_rows
        else:
            r.fetchall.return_value = []
        return r

    conn.execute.side_effect = fake_execute
    with (
        patch("llm.index._engine", return_value=engine),
        patch("llm.index.vectorstore", return_value=store),
        patch("llm.index.embed_texts", return_value=[[0.0] * 3]),
        patch("llm.index.ensure_hnsw"),
        patch("llm.index.enrich_pdf_pages", side_effect=chunks_for) as enrich,
        patch("llm.index.HostPool") as MockPool,
    ):
        MockPool.return_value.check.return_value = ["http://h1:11434"]
        stats = index_refs(
            refs,
            collection="comments",
            cfg=CFG,
            pool=MockPool.return_value,
            force=force,
            sealed=sealed,
            mark_complete=mark_complete,
        )
    return stats, store, conn, enrich


def test_fingerprint_match_skips_without_load():
    loads = []
    stats, store, _, enrich = _run([_ref("fp1", loads=loads)], [("K1", "h-old", "fp1")])
    assert stats["fingerprint_skipped"] == 1 and stats["indexed"] == 0
    assert loads == []
    store.add_embeddings.assert_not_called()
    enrich.assert_not_called()


def test_hash_match_updates_fingerprint_without_embedding():
    h = content_hash(DOC.text)
    stats, store, conn, enrich = _run([_ref("fp2")], [("K1", h, "fp1")])
    assert stats["hash_skipped"] == 1 and stats["indexed"] == 0
    store.add_embeddings.assert_not_called()
    enrich.assert_not_called()
    upd = [
        c
        for c in conn.execute.call_args_list
        if "UPDATE index_state SET fingerprint" in str(c.args[0])
    ]
    assert len(upd) == 1 and upd[0].args[1]["f"] == "fp2"


def test_changed_doc_embeds_and_records_fingerprint():
    stats, store, conn, enrich = _run([_ref("fp2")], [("K1", "stale", "fp1")])
    assert stats["indexed"] == 1 and stats["chunks"] == 1
    store.add_embeddings.assert_called_once()
    enrich.assert_called_once()
    ins = [
        c
        for c in conn.execute.call_args_list
        if "INSERT INTO index_state" in str(c.args[0])
    ]
    assert ins[0].args[1]["f"] == "fp2"


def test_empty_fingerprint_never_matches():
    stats, store, _, _ = _run([_ref("")], [("K1", "stale", "")])
    assert stats["indexed"] == 1


def test_force_ignores_fingerprint_and_hash():
    h = content_hash(DOC.text)
    stats, store, _, _ = _run([_ref("fp1")], [("K1", h, "fp1")], force=True)
    assert stats["indexed"] == 1


def test_load_none_counts_skipped():
    ref = DocRef(
        key="K9", collection="comments", docket="D", fingerprint="x", load=lambda: None
    )
    stats, *_ = _run([ref], [])
    assert stats["skipped"] == 1 and stats["indexed"] == 0


def test_sealed_docket_marked_complete_after_clean_run():
    stats, _, conn, _ = _run([_ref("fp1")], [], sealed={"D": "2026-10-20T00:00:00Z"})
    assert stats["docket_complete"] == 1
    ins = [
        c
        for c in conn.execute.call_args_list
        if "INSERT INTO index_docket_state" in str(c.args[0])
    ]
    assert ins[0].args[1] == {"c": "comments", "d": "D", "s": "2026-10-20T00:00:00Z"}


def test_not_marked_when_mark_complete_false_or_unsealed():
    stats, _, conn, _ = _run([_ref("fp1")], [], sealed={"D": "s"}, mark_complete=False)
    assert stats["docket_complete"] == 0
    stats, _, conn, _ = _run([_ref("fp1")], [], sealed={})
    assert stats["docket_complete"] == 0


def test_sealed_docket_not_marked_when_a_ref_is_unindexable():
    """A docket with an item that has no indexable text is not complete."""
    blank = DocRef(
        key="K9", collection="comments", docket="D", fingerprint="x", load=lambda: None
    )
    stats, _, conn, _ = _run(
        [_ref("fp1"), blank], [], sealed={"D": "2026-10-20T00:00:00Z"}
    )
    assert stats["skipped"] == 1 and stats["indexed"] == 1
    assert stats["docket_complete"] == 0
    assert not [
        c
        for c in conn.execute.call_args_list
        if "INSERT INTO index_docket_state" in str(c.args[0])
    ]


def test_sealed_docket_not_marked_when_a_ref_yields_no_chunks():
    stats, _, conn, _ = _run(
        [_ref("fp1")], [], sealed={"D": "s"}, chunks_for=lambda d, c: []
    )
    assert stats["skipped"] == 1 and stats["docket_complete"] == 0
    assert not [
        c
        for c in conn.execute.call_args_list
        if "INSERT INTO index_docket_state" in str(c.args[0])
    ]