QMDSearch/tests/test_qdrant_admin.py

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