"""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