tests/test_vector_index.py

import os
import numpy as np
import pytest
from unittest.mock import MagicMock, patch


@pytest.fixture
def csv_file(tmp_path):
    p = tmp_path / "index.csv"
    p.write_text(
        "path,token,phrase,word\n"
        "video_dir/hello.mp4,HELLO,,\n"
        "video_dir/thank_you.mp4,THANK-YOU,,\n"
    )
    return str(p)


@pytest.fixture
def chroma_dir(tmp_path):
    return str(tmp_path / "chroma")


def _empty_client():
    col = MagicMock()
    col.count.return_value = 0
    client = MagicMock()
    client.get_or_create_collection.return_value = col
    return client, col


def _full_client():
    col = MagicMock()
    col.count.return_value = 50
    client = MagicMock()
    client.get_or_create_collection.return_value = col
    return client, col


def test_builds_index_when_collection_empty(csv_file, chroma_dir):
    """On first init, upserts all CSV token keys into ChromaDB."""
    client, col = _empty_client()
    model = MagicMock()
    model.encode.return_value = np.array([[0.1, 0.2], [0.3, 0.4]])

    with patch("chromadb.PersistentClient", return_value=client), \
         patch("vector_index.SentenceTransformer", return_value=model):
        from vector_index import VectorIndex
        VectorIndex(csv_path=csv_file, chroma_dir=chroma_dir)

    col.upsert.assert_called_once()
    ids = col.upsert.call_args.kwargs["ids"]
    assert "HELLO" in ids
    assert "THANK-YOU" in ids


def test_skips_build_when_collection_populated(csv_file, chroma_dir):
    """Does not re-encode if collection already has data."""
    client, col = _full_client()
    model = MagicMock()

    with patch("chromadb.PersistentClient", return_value=client), \
         patch("vector_index.SentenceTransformer", return_value=model):
        from vector_index import VectorIndex
        VectorIndex(csv_path=csv_file, chroma_dir=chroma_dir)

    col.upsert.assert_not_called()


def test_query_gloss_returns_above_threshold(csv_file, chroma_dir):
    """Returns (key, score, filename) tuples for results above SEMANTIC_THRESHOLD."""
    client, col = _full_client()
    col.query.return_value = {
        "documents": [["HELP"]],
        "distances": [[0.15]],   # score = 1 - 0.15 = 0.85
        "metadatas": [[{"filename": "help.mp4"}]],
    }
    model = MagicMock()
    model.encode.return_value = np.array([[0.1, 0.2]])

    with patch("chromadb.PersistentClient", return_value=client), \
         patch("vector_index.SentenceTransformer", return_value=model):
        from vector_index import VectorIndex
        idx = VectorIndex(csv_path=csv_file, chroma_dir=chroma_dir)

    with patch.dict(os.environ, {"SEMANTIC_THRESHOLD": "0.6"}):
        results = idx.query_gloss("ASSIST")

    assert len(results) == 1
    key, score, fname = results[0]
    assert key == "HELP"
    assert abs(score - 0.85) < 0.001
    assert fname == "help.mp4"


def test_query_gloss_filters_below_threshold(csv_file, chroma_dir):
    """Returns [] when best cosine score is below SEMANTIC_THRESHOLD."""
    client, col = _full_client()
    col.query.return_value = {
        "documents": [["HELP"]],
        "distances": [[0.7]],    # score = 1 - 0.7 = 0.3 < 0.6
        "metadatas": [[{"filename": "help.mp4"}]],
    }
    model = MagicMock()
    model.encode.return_value = np.array([[0.1, 0.2]])

    with patch("chromadb.PersistentClient", return_value=client), \
         patch("vector_index.SentenceTransformer", return_value=model):
        from vector_index import VectorIndex
        idx = VectorIndex(csv_path=csv_file, chroma_dir=chroma_dir)

    with patch.dict(os.environ, {"SEMANTIC_THRESHOLD": "0.6"}):
        results = idx.query_gloss("XYZZY")

    assert results == []