import tempfile from pathlib import Path from unittest.mock import MagicMock, patch import pytest 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" @patch("companion.rag.search.OllamaEmbedder") def test_search_raises_on_embedder_failure(mock_embedder_cls): mock_embedder = MagicMock() mock_embedder.embed.side_effect = RuntimeError("Connection failed") mock_embedder_cls.return_value = mock_embedder with tempfile.TemporaryDirectory() as tmp: store = VectorStore(uri=tmp, dimensions=4) 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, ) with pytest.raises(RuntimeError, match="Failed to generate embedding"): engine.search("hello")