82 lines
2.7 KiB
Python
82 lines
2.7 KiB
Python
from unittest.mock import MagicMock
|
|
import pytest
|
|
|
|
from app.qdrant_store import (
|
|
ensure_collection,
|
|
upsert_chunks,
|
|
delete_by_path,
|
|
ChunkPoint,
|
|
)
|
|
|
|
|
|
def test_ensure_collection_creates_when_missing():
|
|
fake_client = MagicMock()
|
|
fake_client.collection_exists.return_value = False
|
|
|
|
ensure_collection(fake_client, "rag_test", vector_size=1024)
|
|
|
|
fake_client.create_collection.assert_called_once()
|
|
args, kwargs = fake_client.create_collection.call_args
|
|
assert kwargs["collection_name"] == "rag_test"
|
|
# Payload indexes get created
|
|
assert fake_client.create_payload_index.call_count == 3
|
|
|
|
|
|
def test_ensure_collection_skips_when_exists_with_matching_dim():
|
|
fake_client = MagicMock()
|
|
fake_client.collection_exists.return_value = True
|
|
info = MagicMock()
|
|
info.config.params.vectors.size = 1024
|
|
fake_client.get_collection.return_value = info
|
|
|
|
ensure_collection(fake_client, "rag_test", vector_size=1024)
|
|
|
|
fake_client.create_collection.assert_not_called()
|
|
|
|
|
|
def test_ensure_collection_raises_on_dim_mismatch():
|
|
fake_client = MagicMock()
|
|
fake_client.collection_exists.return_value = True
|
|
info = MagicMock()
|
|
info.config.params.vectors.size = 768
|
|
fake_client.get_collection.return_value = info
|
|
|
|
with pytest.raises(RuntimeError, match="dimension mismatch"):
|
|
ensure_collection(fake_client, "rag_test", vector_size=1024)
|
|
|
|
|
|
def test_upsert_chunks_calls_client_upsert():
|
|
fake_client = MagicMock()
|
|
points = [
|
|
ChunkPoint(vector=[0.1] * 4, payload={"file_path": "a", "chunk_index": 0}),
|
|
ChunkPoint(vector=[0.2] * 4, payload={"file_path": "a", "chunk_index": 1}),
|
|
]
|
|
|
|
upsert_chunks(fake_client, "rag_test", points)
|
|
|
|
fake_client.upsert.assert_called_once()
|
|
kwargs = fake_client.upsert.call_args.kwargs
|
|
assert kwargs["collection_name"] == "rag_test"
|
|
produced = kwargs["points"]
|
|
assert len(produced) == 2
|
|
assert produced[0].vector == [0.1] * 4
|
|
assert produced[0].payload == {"file_path": "a", "chunk_index": 0}
|
|
assert produced[1].vector == [0.2] * 4
|
|
assert produced[1].payload == {"file_path": "a", "chunk_index": 1}
|
|
# ids are UUID strings, distinct
|
|
assert isinstance(produced[0].id, str)
|
|
assert produced[0].id != produced[1].id
|
|
|
|
|
|
def test_delete_by_path_uses_filter():
|
|
fake_client = MagicMock()
|
|
delete_by_path(fake_client, "rag_test", "Documents/x.pdf")
|
|
|
|
fake_client.delete.assert_called_once()
|
|
kwargs = fake_client.delete.call_args.kwargs
|
|
assert kwargs["collection_name"] == "rag_test"
|
|
# The filter should target file_path
|
|
selector = kwargs["points_selector"]
|
|
# Inspect the FilterSelector → Filter → must → FieldCondition
|
|
assert selector.filter.must[0].key == "file_path"
|