298 lines
9.1 KiB
Python
298 lines
9.1 KiB
Python
"""M6 Qdrant 向量存储模块单元测试。"""
|
|
|
|
|
|
|
|
from vectorstore.client import (
|
|
DEFAULT_COLLECTION,
|
|
VectorStore,
|
|
make_qdrant_client,
|
|
url_hash_to_uuid,
|
|
)
|
|
from vectorstore.models import CollectionInfo, SearchFilter, SearchResult
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# url_hash_to_uuid
|
|
# --------------------------------------------------------------------------- #
|
|
|
|
|
|
class TestUrlHashToUuid:
|
|
"""url_hash_to_uuid 函数测试。"""
|
|
|
|
def test_deterministic(self):
|
|
h = "abc1234567890000"
|
|
assert url_hash_to_uuid(h) == url_hash_to_uuid(h)
|
|
|
|
def test_different_hash_different_uuid(self):
|
|
assert url_hash_to_uuid("aaa1111111111111") != url_hash_to_uuid("bbb2222222222222")
|
|
|
|
def test_valid_uuid_format(self):
|
|
import uuid
|
|
result = url_hash_to_uuid("abc1234567890000")
|
|
uuid.UUID(result) # 应不抛异常
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# make_qdrant_client
|
|
# --------------------------------------------------------------------------- #
|
|
|
|
|
|
class TestMakeQdrantClient:
|
|
"""make_qdrant_client 测试。"""
|
|
|
|
def test_memory_mode(self):
|
|
client = make_qdrant_client(memory=True)
|
|
assert client is not None
|
|
client.close()
|
|
|
|
def test_custom_path(self, tmp_path):
|
|
path = str(tmp_path / "qdrant_test")
|
|
client = make_qdrant_client(path=path)
|
|
assert client is not None
|
|
client.close()
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# VectorStore — Collection 管理
|
|
# --------------------------------------------------------------------------- #
|
|
|
|
|
|
class TestVectorStoreInit:
|
|
"""VectorStore Collection 初始化测试。"""
|
|
|
|
def test_init_collection_creates(self):
|
|
client = make_qdrant_client(memory=True)
|
|
store = VectorStore(client)
|
|
try:
|
|
store.init_collection()
|
|
info = store.info()
|
|
assert info.exists
|
|
assert info.name == DEFAULT_COLLECTION
|
|
finally:
|
|
store.close()
|
|
|
|
def test_init_idempotent(self):
|
|
client = make_qdrant_client(memory=True)
|
|
store = VectorStore(client)
|
|
try:
|
|
store.init_collection()
|
|
store.init_collection() # 第二次不应报错
|
|
assert store.info().exists
|
|
finally:
|
|
store.close()
|
|
|
|
def test_recreate(self):
|
|
client = make_qdrant_client(memory=True)
|
|
store = VectorStore(client)
|
|
try:
|
|
store.init_collection()
|
|
store.init_collection(recreate=True)
|
|
assert store.info().exists
|
|
finally:
|
|
store.close()
|
|
|
|
def test_info_nonexistent(self):
|
|
client = make_qdrant_client(memory=True)
|
|
store = VectorStore(client, collection_name="nonexistent_test")
|
|
try:
|
|
info = store.info()
|
|
assert not info.exists
|
|
finally:
|
|
store.close()
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# VectorStore — upsert + query
|
|
# --------------------------------------------------------------------------- #
|
|
|
|
|
|
class TestVectorStoreUpsert:
|
|
"""VectorStore upsert 测试。"""
|
|
|
|
def test_upsert_and_count(self):
|
|
client = make_qdrant_client(memory=True)
|
|
store = VectorStore(client)
|
|
try:
|
|
store.init_collection()
|
|
points = [
|
|
{
|
|
"id": "hash0000000000001",
|
|
"vector": [0.1] * 1024,
|
|
"payload": {
|
|
"title": "Test Article",
|
|
"title_zh": "测试文章",
|
|
"url": "https://example.com/1",
|
|
"source_id": "reuters",
|
|
},
|
|
},
|
|
{
|
|
"id": "hash0000000000002",
|
|
"vector": [0.2] * 1024,
|
|
"payload": {
|
|
"title": "Another Article",
|
|
"title_zh": "另一篇文章",
|
|
"url": "https://example.com/2",
|
|
"source_id": "cnbc",
|
|
},
|
|
},
|
|
]
|
|
count = store.upsert(points)
|
|
assert count == 2
|
|
assert store.count() == 2
|
|
finally:
|
|
store.close()
|
|
|
|
def test_upsert_idempotent(self):
|
|
client = make_qdrant_client(memory=True)
|
|
store = VectorStore(client)
|
|
try:
|
|
store.init_collection()
|
|
points = [{
|
|
"id": "hash0000000000001",
|
|
"vector": [0.1] * 1024,
|
|
"payload": {"title": "Original"},
|
|
}]
|
|
store.upsert(points)
|
|
|
|
# 同一 id 第二次写入（更新）
|
|
points2 = [{
|
|
"id": "hash0000000000001",
|
|
"vector": [0.9] * 1024,
|
|
"payload": {"title": "Updated"},
|
|
}]
|
|
store.upsert(points2)
|
|
|
|
assert store.count() == 1 # 不应增加
|
|
finally:
|
|
store.close()
|
|
|
|
|
|
class TestVectorStoreQuery:
|
|
"""VectorStore query 测试。"""
|
|
|
|
def test_query_returns_results(self):
|
|
client = make_qdrant_client(memory=True)
|
|
store = VectorStore(client)
|
|
try:
|
|
store.init_collection()
|
|
# 写入 3 条
|
|
for i in range(3):
|
|
store.upsert([{
|
|
"id": f"hash{i:016d}",
|
|
"vector": [float(i) / 10] * 1024,
|
|
"payload": {
|
|
"title": f"Article {i}",
|
|
"title_zh": f"文章 {i}",
|
|
"url": f"https://example.com/{i}",
|
|
"source_id": "reuters",
|
|
"events": [],
|
|
},
|
|
}])
|
|
|
|
# 查询
|
|
query_vec = [0.15] * 1024 # 接近 hash0 和 hash1
|
|
results = store.query(query_vec, top_k=2)
|
|
assert len(results) == 2
|
|
assert results[0].score > 0 # 有相似度分数
|
|
finally:
|
|
store.close()
|
|
|
|
def test_query_with_filter(self):
|
|
client = make_qdrant_client(memory=True)
|
|
store = VectorStore(client)
|
|
try:
|
|
store.init_collection()
|
|
for i in range(5):
|
|
store.upsert([{
|
|
"id": f"hash{i:016d}",
|
|
"vector": [0.5] * 1024,
|
|
"payload": {
|
|
"title": f"A{i}",
|
|
"title_zh": f"文{i}",
|
|
"url": f"https://x.com/{i}",
|
|
"source_id": "reuters" if i < 3 else "cnbc",
|
|
"events": [],
|
|
},
|
|
}])
|
|
|
|
# 过滤只查 reuters
|
|
sf = SearchFilter(source_id="reuters")
|
|
results = store.query([0.5] * 1024, top_k=10, search_filter=sf)
|
|
assert all(r.source_id == "reuters" for r in results)
|
|
finally:
|
|
store.close()
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# SearchFilter / SearchResult 模型
|
|
# --------------------------------------------------------------------------- #
|
|
|
|
|
|
class TestSearchFilter:
|
|
"""SearchFilter 模型测试。"""
|
|
|
|
def test_empty_filter(self):
|
|
f = SearchFilter()
|
|
assert f.source_id is None
|
|
|
|
def test_source_id_filter(self):
|
|
f = SearchFilter(source_id="reuters")
|
|
assert f.source_id == "reuters"
|
|
|
|
def test_multi_condition(self):
|
|
f = SearchFilter(
|
|
source_ids=["reuters", "cnbc"],
|
|
sentiment="positive",
|
|
importance_min=3,
|
|
publish_date_from="2026-06-01",
|
|
publish_date_to="2026-06-30",
|
|
)
|
|
assert f.sentiment == "positive"
|
|
assert f.importance_min == 3
|
|
|
|
|
|
class TestSearchResult:
|
|
"""SearchResult 模型测试。"""
|
|
|
|
def test_basic(self):
|
|
r = SearchResult(
|
|
url_hash="abc",
|
|
score=0.95,
|
|
title="Fed Holds Rates",
|
|
title_zh="美联储维持利率",
|
|
url="https://example.com/1",
|
|
source_id="reuters",
|
|
)
|
|
assert r.score == 0.95
|
|
assert "美联储" in r.short_summary()
|
|
|
|
def test_with_events(self):
|
|
r = SearchResult(
|
|
url_hash="abc",
|
|
score=0.88,
|
|
title="Apple Earnings",
|
|
title_zh="苹果财报",
|
|
source_id="reuters",
|
|
events=[{
|
|
"event_type": "财报披露",
|
|
"stock_codes": ["AAPL"],
|
|
"sentiment": "positive",
|
|
"importance": 4,
|
|
"summary_zh": "苹果财报超预期",
|
|
}],
|
|
)
|
|
assert "AAPL" in r.short_summary()
|
|
|
|
|
|
class TestCollectionInfo:
|
|
"""CollectionInfo 模型测试。"""
|
|
|
|
def test_basic(self):
|
|
info = CollectionInfo(name="test", exists=True, vectors_count=42)
|
|
assert info.exists
|
|
assert info.vectors_count == 42
|
|
|
|
def test_defaults(self):
|
|
info = CollectionInfo(name="empty", exists=False)
|
|
assert info.vectors_count == 0
|
|
assert info.indexed_vectors_count is None
|