35 lines
1.1 KiB
Python
35 lines
1.1 KiB
Python
import tempfile
|
|
from pathlib import Path
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
from companion.rag.search import SearchEngine
|
|
from companion.rag.vector_store import VectorStore
|
|
|
|
|
|
@patch("companion.rag.search.OllamaEmbedder")
|
|
def test_search_returns_results(mock_embedder_cls):
|
|
mock_embedder = MagicMock()
|
|
mock_embedder.embed.return_value = [[1.0, 0.0, 0.0, 0.0]]
|
|
mock_embedder_cls.return_value = mock_embedder
|
|
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
store = VectorStore(uri=tmp, dimensions=4)
|
|
store.upsert(
|
|
ids=["a"],
|
|
texts=["hello world"],
|
|
embeddings=[[1.0, 0.0, 0.0, 0.0]],
|
|
metadatas=[{"source_file": "a.md", "source_directory": "docs"}],
|
|
)
|
|
engine = SearchEngine(
|
|
vector_store=store,
|
|
embedder_base_url="http://localhost:11434",
|
|
embedder_model="dummy",
|
|
embedder_batch_size=32,
|
|
default_top_k=5,
|
|
similarity_threshold=0.0,
|
|
hybrid_search_enabled=False,
|
|
)
|
|
results = engine.search("hello")
|
|
assert len(results) == 1
|
|
assert results[0]["source_file"] == "a.md"
|