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 == []