tests/test_store.py

"""Unit tests for src/store.py pure functions — no models, no Qdrant instance."""
import src.store as store


def _unit(cid, parent, chapter="Chapter 1: X", **kw):
    return {"id": cid, "parent_id": parent, "chapter": chapter, "section": "",
            "text": f"text-{cid}", "is_parent": kw.get("is_parent", False),
            "type": "child", "source": "s", **kw}


def test_contextualize_prepends_location():
    assert store.contextualize("Ch", "Sec", "body") == "[Ch > Sec]\nbody"
    assert store.contextualize("Ch", "", "body") == "[Ch]\nbody"
    assert store.contextualize("", "", "body") == "body"


def test_retrieval_units_selection():
    chunks = [
        # parent with children -> represented by them, not itself
        _unit("p1", None, is_parent=True),
        _unit("c1", "p1"), _unit("c2", "p1"),
        # childless parent -> its own unit
        _unit("p2", None, is_parent=True),
        # front/back matter (empty chapter) -> dropped
        _unit("p3", None, is_parent=True, chapter=""),
    ]
    units = store.retrieval_units(chunks)
    ids = [u["id"] for u in units]
    assert ids == ["c1", "c2", "p2"]
    assert all(u["parent_text"] == "text-p1" for u in units[:2])  # children answer with parent
    assert units[2]["parent_text"] == "text-p2"


def test_dedupe_parents_keeps_first_per_parent():
    hits = [{"parent_id": "a", "v": 1}, {"parent_id": "b", "v": 2}, {"parent_id": "a", "v": 3}]
    out = store._dedupe_parents(hits)
    assert [h["v"] for h in out] == [1, 2]


def test_sparse_vec_roundtrip():
    v = store._sparse_vec({3: 0.5, 7: 1.0})
    assert list(v.indices) == [3, 7]
    assert list(v.values) == [0.5, 1.0]


def test_rerank_empty_hits():
    assert store.rerank("q", [], top_k=5) == []


def test_search_multi_round_robin_merge_dedupes_and_caps(monkeypatch):
    # Two lanes; lane hits are pre-ranked. parent p2 appears in both lanes.
    lanes = {
        "qA": [{"parent_id": "p1"}, {"parent_id": "p2"}, {"parent_id": "p3"}],
        "qB": [{"parent_id": "p2"}, {"parent_id": "p4"}, {"parent_id": "p5"}],
    }
    monkeypatch.setattr(store, "_hybrid_candidates", lambda c, q, cand, name: lanes[q])
    monkeypatch.setattr(store, "rerank", lambda q, hits, top_k: hits[:top_k])
    out = store.search_multi(None, ["qA", "qB"], top_k=4)
    # round-robin by rank: p1 (A1), p2 (B1), then A2 (=p2, dup skipped), p4 (B2), p3 (A3)
    assert [h["parent_id"] for h in out] == ["p1", "p2", "p4", "p3"]
    assert len(out) == 4  # capped at top_k


def test_search_multi_guarantees_every_lane_representation(monkeypatch):
    lanes = {
        "poisson": [{"parent_id": f"po{i}"} for i in range(5)],
        "geometric": [{"parent_id": f"ge{i}"} for i in range(5)],
    }
    monkeypatch.setattr(store, "_hybrid_candidates", lambda c, q, cand, name: lanes[q])
    monkeypatch.setattr(store, "rerank", lambda q, hits, top_k: hits[:top_k])
    out = store.search_multi(None, ["poisson", "geometric"], top_k=4)
    pids = [h["parent_id"] for h in out]
    assert any(p.startswith("po") for p in pids) and any(p.startswith("ge") for p in pids)