pagepiper/tests/test_colbert_index.py

# tests/test_colbert_index.py
"""Tests for app.services.colbert_index.

pylate is NOT installed in the dev/test env by design (see cf-sysadmin skill's
"Known Gotchas" — installing it directly into the shared `cf` conda env broke
several other services' pinned torch/transformers versions on 2026-07-10).
These tests inject fake `pylate`/`pylate.models`/`pylate.indexes`/`pylate.retrieve`
modules via sys.modules so ColBERTIndex's lazy imports resolve to mocks without
pylate ever needing to be installed here.
"""
from __future__ import annotations

import sqlite3
import sys
import types
from pathlib import Path
from unittest.mock import MagicMock

import pytest

from app.services.colbert_index import ColBERTIndex


@pytest.fixture
def fake_pylate(monkeypatch):
    fake_models_mod = types.ModuleType("pylate.models")
    fake_indexes_mod = types.ModuleType("pylate.indexes")
    fake_retrieve_mod = types.ModuleType("pylate.retrieve")
    fake_pylate_mod = types.ModuleType("pylate")
    fake_pylate_mod.models = fake_models_mod
    fake_pylate_mod.indexes = fake_indexes_mod
    fake_pylate_mod.retrieve = fake_retrieve_mod

    mock_model = MagicMock()
    mock_model.encode.return_value = [[0.1, 0.2], [0.3, 0.4]]
    fake_models_mod.ColBERT = MagicMock(return_value=mock_model)

    mock_index = MagicMock()
    fake_indexes_mod.Voyager = MagicMock(return_value=mock_index)

    mock_retriever = MagicMock()
    fake_retrieve_mod.ColBERT = MagicMock(return_value=mock_retriever)

    monkeypatch.setitem(sys.modules, "pylate", fake_pylate_mod)
    monkeypatch.setitem(sys.modules, "pylate.models", fake_models_mod)
    monkeypatch.setitem(sys.modules, "pylate.indexes", fake_indexes_mod)
    monkeypatch.setitem(sys.modules, "pylate.retrieve", fake_retrieve_mod)

    return types.SimpleNamespace(
        model_cls=fake_models_mod.ColBERT,
        model=mock_model,
        index_cls=fake_indexes_mod.Voyager,
        index=mock_index,
        retriever_cls=fake_retrieve_mod.ColBERT,
        retriever=mock_retriever,
    )


@pytest.fixture
def seeded_db(tmp_path) -> str:
    db_path = str(tmp_path / "test.db")
    schema = Path("migrations/001_initial_schema.sql").read_text()
    conn = sqlite3.connect(db_path)
    conn.executescript(schema)
    conn.execute(
        "INSERT INTO documents(id, title, file_path, status) VALUES ('d1','Test','test.pdf','ready')"
    )
    conn.execute(
        "INSERT INTO page_chunks(id, doc_id, page_number, text, source, word_count) "
        "VALUES ('c1','d1',1,'Setting the IP on the AVC-X','text',6)"
    )
    conn.execute(
        "INSERT INTO page_chunks(id, doc_id, page_number, text, source, word_count) "
        "VALUES ('c2','d1',2,'Filter cartridge replacement steps','text',5)"
    )
    conn.commit()
    conn.close()
    return db_path


def test_ensure_fresh_builds_index_from_sqlite(fake_pylate, seeded_db, tmp_path):
    idx = ColBERTIndex(index_dir=str(tmp_path / "colbert_index"))
    idx.ensure_fresh(seeded_db)

    fake_pylate.model.encode.assert_called_once()
    call_args = fake_pylate.model.encode.call_args
    assert set(call_args[0][0]) == {"Setting the IP on the AVC-X", "Filter cartridge replacement steps"}
    assert call_args[1]["is_query"] is False

    fake_pylate.index.add_documents.assert_called_once()
    add_kwargs = fake_pylate.index.add_documents.call_args[1]
    assert set(add_kwargs["documents_ids"]) == {"c1", "c2"}


def test_ensure_fresh_skips_rebuild_when_not_dirty(fake_pylate, seeded_db, tmp_path):
    idx = ColBERTIndex(index_dir=str(tmp_path / "colbert_index"))
    idx.ensure_fresh(seeded_db)
    idx.ensure_fresh(seeded_db)

    fake_pylate.model.encode.assert_called_once()


def test_mark_dirty_forces_rebuild(fake_pylate, seeded_db, tmp_path):
    idx = ColBERTIndex(index_dir=str(tmp_path / "colbert_index"))
    idx.ensure_fresh(seeded_db)
    idx.mark_dirty()
    idx.ensure_fresh(seeded_db)

    assert fake_pylate.model.encode.call_count == 2


def test_ensure_fresh_with_empty_corpus_leaves_index_none(fake_pylate, tmp_path):
    db_path = str(tmp_path / "empty.db")
    schema = Path("migrations/001_initial_schema.sql").read_text()
    conn = sqlite3.connect(db_path)
    conn.executescript(schema)
    conn.commit()
    conn.close()

    idx = ColBERTIndex(index_dir=str(tmp_path / "colbert_index"))
    idx.ensure_fresh(db_path)

    fake_pylate.index_cls.assert_not_called()
    assert idx.query("anything") == []


def test_query_maps_results_back_to_chunks(fake_pylate, seeded_db, tmp_path):
    idx = ColBERTIndex(index_dir=str(tmp_path / "colbert_index"))
    idx.ensure_fresh(seeded_db)

    fake_pylate.retriever.retrieve.return_value = [
        [{"id": "c1", "score": 13.5}, {"id": "c2", "score": 9.2}]
    ]

    results = idx.query("how do I set the IP on the AVC-X", top_k=10)

    assert len(results) == 2
    assert results[0]["chunk_id"] == "c1"
    assert results[0]["doc_id"] == "d1"
    assert results[0]["score"] == 13.5
    assert results[1]["chunk_id"] == "c2"


def test_query_filters_by_doc_ids(fake_pylate, seeded_db, tmp_path):
    idx = ColBERTIndex(index_dir=str(tmp_path / "colbert_index"))
    idx.ensure_fresh(seeded_db)

    fake_pylate.retriever.retrieve.return_value = [
        [{"id": "c1", "score": 13.5}, {"id": "c2", "score": 9.2}]
    ]

    results = idx.query("query", top_k=10, doc_ids=["other-doc"])

    assert results == []


def test_query_respects_top_k(fake_pylate, seeded_db, tmp_path):
    idx = ColBERTIndex(index_dir=str(tmp_path / "colbert_index"))
    idx.ensure_fresh(seeded_db)

    fake_pylate.retriever.retrieve.return_value = [
        [{"id": "c1", "score": 13.5}, {"id": "c2", "score": 9.2}]
    ]

    results = idx.query("query", top_k=1)

    assert len(results) == 1
    assert results[0]["chunk_id"] == "c1"