Spaces:
Sleeping
Sleeping
| import sys | |
| from pathlib import Path | |
| sys.path.insert(0, str(Path(__file__).parent.parent / "src")) | |
| import pytest | |
| import numpy as np | |
| from unittest.mock import Mock, patch, MagicMock | |
| from langchain_core.embeddings import Embeddings | |
| from config import EMBEDDING_DIMENSION | |
| class MockEmbedder(Embeddings): | |
| """Mock embedder implementing LangChain Embeddings interface.""" | |
| def embed_documents(self, texts): | |
| return [[0.1] * EMBEDDING_DIMENSION for _ in texts] | |
| def embed_query(self, text): | |
| return [0.1] * EMBEDDING_DIMENSION | |
| class TestVectorStore: | |
| """Tests for VectorStore class.""" | |
| def test_init_memory_mode(self): | |
| """Should initialize in memory mode.""" | |
| from vector_store import VectorStore, reset_vector_store | |
| reset_vector_store() | |
| store = VectorStore(use_memory=True, embedder=MockEmbedder()) | |
| assert store.use_memory is True | |
| assert store._client is not None | |
| def test_collection_created_on_init(self): | |
| """Collection should be created during init.""" | |
| from vector_store import VectorStore, reset_vector_store | |
| reset_vector_store() | |
| store = VectorStore(use_memory=True, embedder=MockEmbedder()) | |
| collections = store._client.get_collections().collections | |
| names = [c.name for c in collections] | |
| assert store.collection_name in names | |
| def test_add_documents(self): | |
| """Should add documents to store.""" | |
| from vector_store import VectorStore, reset_vector_store | |
| reset_vector_store() | |
| store = VectorStore(use_memory=True, embedder=MockEmbedder()) | |
| texts = ["Document one", "Document two"] | |
| metadatas = [{"source": "a.pdf"}, {"source": "b.pdf"}] | |
| ids = store.add_documents(texts, metadatas) | |
| assert len(ids) == 2 | |
| def test_add_documents_empty_list(self): | |
| """Empty list should return empty list.""" | |
| from vector_store import VectorStore, reset_vector_store | |
| reset_vector_store() | |
| store = VectorStore(use_memory=True, embedder=MockEmbedder()) | |
| ids = store.add_documents([]) | |
| assert ids == [] | |
| def test_add_documents_without_metadata(self): | |
| """Should work without metadata.""" | |
| from vector_store import VectorStore, reset_vector_store | |
| reset_vector_store() | |
| store = VectorStore(use_memory=True, embedder=MockEmbedder()) | |
| ids = store.add_documents(["Test document"]) | |
| assert len(ids) == 1 | |
| def test_search_returns_formatted_results(self): | |
| """Search should return properly formatted results.""" | |
| from vector_store import VectorStore, reset_vector_store | |
| reset_vector_store() | |
| store = VectorStore(use_memory=True, embedder=MockEmbedder()) | |
| store.add_documents( | |
| ["Test content here"], | |
| [{"source": "test.pdf", "chunk_index": 0, "page_number": 1}] | |
| ) | |
| results = store.search("test", top_k=1) | |
| assert len(results) == 1 | |
| assert "score" in results[0] | |
| assert "text" in results[0] | |
| assert "source" in results[0] | |
| assert "chunk_index" in results[0] | |
| assert "page_number" in results[0] | |
| def test_get_collection_stats(self): | |
| """Should return collection statistics.""" | |
| from vector_store import VectorStore, reset_vector_store | |
| reset_vector_store() | |
| store = VectorStore(use_memory=True, embedder=MockEmbedder()) | |
| store.clear_collection() | |
| store.add_documents(["Doc 1", "Doc 2"]) | |
| stats = store.get_collection_stats() | |
| assert stats["name"] == store.collection_name | |
| assert stats["points_count"] == 2 | |
| def test_clear_collection(self): | |
| """Should clear all documents.""" | |
| from vector_store import VectorStore, reset_vector_store | |
| reset_vector_store() | |
| store = VectorStore(use_memory=True, embedder=MockEmbedder()) | |
| store.add_documents(["Doc 1", "Doc 2"]) | |
| store.clear_collection() | |
| stats = store.get_collection_stats() | |
| assert stats["points_count"] == 0 | |
| def test_collection_exists_false_when_empty(self): | |
| """Should return False for empty collection.""" | |
| from vector_store import VectorStore, reset_vector_store | |
| reset_vector_store() | |
| store = VectorStore(use_memory=True, embedder=MockEmbedder()) | |
| assert store.collection_exists() is False | |
| def test_collection_exists_true_with_docs(self): | |
| """Should return True when documents exist.""" | |
| from vector_store import VectorStore, reset_vector_store | |
| reset_vector_store() | |
| store = VectorStore(use_memory=True, embedder=MockEmbedder()) | |
| store.add_documents(["Test"]) | |
| assert store.collection_exists() is True | |
| def test_metadata_preserved_on_retrieval(self): | |
| """Metadata should be preserved when retrieving documents.""" | |
| from vector_store import VectorStore, reset_vector_store | |
| reset_vector_store() | |
| store = VectorStore(use_memory=True, embedder=MockEmbedder()) | |
| metadata = { | |
| "source": "report.pdf", | |
| "chunk_index": 5, | |
| "page_number": 3, | |
| "custom_field": "custom_value" | |
| } | |
| store.add_documents(["Important content"], [metadata]) | |
| results = store.search("important", top_k=1) | |
| assert results[0]["source"] == "report.pdf" | |
| assert results[0]["chunk_index"] == 5 | |
| assert results[0]["page_number"] == 3 | |
| class TestSingleton: | |
| """Tests for singleton pattern.""" | |
| def test_get_vector_store_returns_same_instance(self): | |
| """get_vector_store should return same instance.""" | |
| from vector_store import reset_vector_store | |
| reset_vector_store() | |
| with patch("vector_store.get_embedder", return_value=MockEmbedder()): | |
| from vector_store import get_vector_store | |
| instance1 = get_vector_store() | |
| instance2 = get_vector_store() | |
| assert instance1 is instance2 | |
| def test_reset_vector_store_clears_instance(self): | |
| """reset_vector_store should clear singleton.""" | |
| from vector_store import reset_vector_store | |
| reset_vector_store() | |
| with patch("vector_store.get_embedder", return_value=MockEmbedder()): | |
| with patch("vector_store.USE_MEMORY_MODE", True): | |
| from vector_store import get_vector_store, VectorStore | |
| instance1 = VectorStore(use_memory=True, embedder=MockEmbedder()) | |
| instance2 = VectorStore(use_memory=True, embedder=MockEmbedder()) | |
| # They are different instances when created directly | |
| assert instance1 is not instance2 | |