"""QdrantService 管理操作测试 使用 AsyncQdrantClient(location=":memory:") 本地模式,无需 Docker。 """ import pytest from qdrant_client import AsyncQdrantClient from app.config import settings from app.services.qdrant import ( ALL_COLLECTIONS, COLLECTION_CHUNKS, COLLECTION_L1, COLLECTION_L2, COLLECTION_L3, QdrantService, ) DIM = settings.embedding_dimension def _dense(seed: float) -> list[float]: """构造确定性 dense 向量:前 4 维取特征值,便于区分不同文档""" vec = [0.0] * DIM vec[0] = seed vec[1] = 1.0 - seed vec[2] = seed * 0.5 vec[3] = 0.1 return vec @pytest.fixture async def service() -> QdrantService: svc = QdrantService(client=AsyncQdrantClient(location=":memory:")) await svc.ensure_collections() return svc async def _seed_doc(service: QdrantService, doc_id: str, chunk_count: int = 2) -> None: """写入一篇完整文档:L1 一条 + L2/L3 各 2 节点 + chunk_count 个 chunk""" await service.upsert_l1( doc_id=doc_id, title=f"标题-{doc_id}", summary=f"总结-{doc_id}", category="tech", tags=["t"], dense_vector=_dense(0.9), ) nodes = [ { "doc_id": doc_id, "section_path": str(i), "text": f"节点{i}-{doc_id}", "category": "tech", "tags": ["t"], "dense_vector": _dense(0.8 - i * 0.1), } for i in range(1, 3) ] await service.upsert_nodes(COLLECTION_L2, nodes) await service.upsert_nodes(COLLECTION_L3, nodes) chunks = [ { "doc_id": doc_id, "chunk_index": i, "text": f"chunk{i}-{doc_id}", "section_path": "1", "title": f"标题-{doc_id}", "category": "tech", "tags": ["t"], "dense_vector": _dense(0.5 + i * 0.1), } for i in range(chunk_count) ] await service.upsert_chunks(chunks) async def test_count_empty_and_after_upsert(service: QdrantService) -> None: for collection in ALL_COLLECTIONS: assert await service.count(collection) == 0 await _seed_doc(service, "doc-count", chunk_count=3) assert await service.count(COLLECTION_L1) == 1 assert await service.count(COLLECTION_L2) == 2 assert await service.count(COLLECTION_L3) == 2 assert await service.count(COLLECTION_CHUNKS) == 3 async def test_scroll_l1_empty(service: QdrantService) -> None: items, next_offset = await service.scroll_l1() assert items == [] assert next_offset is None async def test_scroll_l1_pagination(service: QdrantService) -> None: """写入 3 篇文档,limit=2 翻页应取全且无重复""" for i in range(3): await service.upsert_l1( doc_id=f"doc-{i}", title=f"标题{i}", summary=f"总结{i}", category="tech", tags=["x"], dense_vector=_dense(0.5 + i * 0.1), ) seen: list[dict] = [] offset: str | None = None pages = 0 while True: items, offset = await service.scroll_l1(limit=2, offset=offset) seen.extend(items) pages += 1 if offset is None: break assert pages <= 3 # 防止游标异常导致死循环 assert pages == 2 assert len(seen) == 3 doc_ids = [item["doc_id"] for item in seen] assert len(set(doc_ids)) == 3 # 无重复 assert set(doc_ids) == {"doc-0", "doc-1", "doc-2"} # item 字段完整,summary 映射自 payload text for item in seen: assert set(item.keys()) == {"doc_id", "title", "category", "tags", "summary"} i = int(item["doc_id"].rsplit("-", 1)[1]) assert item["title"] == f"标题{i}" assert item["summary"] == f"总结{i}" assert item["category"] == "tech" assert item["tags"] == ["x"] async def test_get_doc_detail_not_found(service: QdrantService) -> None: assert await service.get_doc_detail("doc-missing") is None async def test_get_doc_detail(service: QdrantService) -> None: await _seed_doc(service, "doc-detail", chunk_count=3) # 干扰数据:不应混入结果 await _seed_doc(service, "doc-other", chunk_count=1) detail = await service.get_doc_detail("doc-detail") assert detail is not None l1 = detail["l1"] assert l1["doc_id"] == "doc-detail" assert l1["title"] == "标题-doc-detail" assert l1["text"] == "总结-doc-detail" assert l1["category"] == "tech" assert len(detail["l2_nodes"]) == 2 assert len(detail["l3_nodes"]) == 2 for nodes in (detail["l2_nodes"], detail["l3_nodes"]): assert {n["section_path"] for n in nodes} == {"1", "2"} for node in nodes: assert node["doc_id"] == "doc-detail" assert "text" in node assert detail["chunks_count"] == 3 async def test_delete_by_doc_id(service: QdrantService) -> None: await _seed_doc(service, "doc-del", chunk_count=2) await _seed_doc(service, "doc-keep", chunk_count=1) deleted = await service.delete_by_doc_id("doc-del") assert deleted == { COLLECTION_L1: 1, COLLECTION_L2: 2, COLLECTION_L3: 2, COLLECTION_CHUNKS: 2, } # 四层该 doc 的点全部清空 assert await service.get_doc_detail("doc-del") is None doc_filter = QdrantService.build_filter(doc_ids=["doc-del"]) for collection in ALL_COLLECTIONS: result = await service.client.count(collection, count_filter=doc_filter, exact=True) assert result.count == 0 # 不影响其他 doc 的数据 keep_detail = await service.get_doc_detail("doc-keep") assert keep_detail is not None assert keep_detail["chunks_count"] == 1 assert len(keep_detail["l2_nodes"]) == 2 async def test_delete_by_doc_id_nonexistent(service: QdrantService) -> None: """删除不存在的 doc_id:各集合返回 0,已有数据不受影响(幂等)""" await _seed_doc(service, "doc-alive", chunk_count=1) deleted = await service.delete_by_doc_id("doc-missing") assert deleted == {name: 0 for name in ALL_COLLECTIONS} assert await service.count(COLLECTION_L1) == 1 assert await service.count(COLLECTION_CHUNKS) == 1