Spaces:
Sleeping
Sleeping
File size: 7,038 Bytes
ba2adac 8853fb6 ba2adac 8853fb6 ba2adac 8853fb6 ba2adac 79d4fd5 ba2adac 8853fb6 ba2adac 79d4fd5 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 | 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
|