116 lines
4.6 KiB
Python
116 lines
4.6 KiB
Python
"""T21-2 milvus_service 单测：真 Milvus Lite（tempfile 临时库，module 级共享）。
|
||
|
||
不 mock：Milvus Lite 本机实测可用（0904 交接 + 0907 复测），真库测试同时
|
||
覆盖 schema 定义 / filter 表达式 / 溯源字段回传的正确性，比 mock 更有判定力。
|
||
|
||
覆盖：
|
||
1. ensure_collection 幂等（两次调用不报错）；
|
||
2. insert_chunks upsert 语义（同 id 重跑覆盖不重复）；
|
||
3. search_kb 相似度排序 + 溯源字段完整；
|
||
4. 合规过滤：effective_date 未生效块必须被排除（§2.3）；
|
||
5. 标量过滤 product_id / doc_type；
|
||
6. 空写入短路。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import tempfile
|
||
import uuid
|
||
from pathlib import Path
|
||
|
||
import pytest
|
||
|
||
from app.service import milvus_service as ms
|
||
|
||
DIM = 16 # 测试维度（与真 embed_dim=1024 解耦，向量手工构造）
|
||
|
||
|
||
@pytest.fixture(scope="module")
|
||
def client():
|
||
"""module 级 Lite 实例（临时目录随机文件，跑完即弃）。"""
|
||
tmp = Path(tempfile.mkdtemp()) / f"{uuid.uuid4().hex}.db"
|
||
c = ms.MilvusClient(uri=str(tmp))
|
||
ms.ensure_collection(c, dim=DIM)
|
||
yield c
|
||
c.close()
|
||
|
||
|
||
def _row(chunk_id: str, text: str, *, product_id="PROD-TEST", doc_type="fee",
|
||
effective_date="2026-01-01", version="v1.0", doc_id="DOC-001") -> dict:
|
||
"""构造一条测试 chunk（向量按文本首字符微调制造差异）。"""
|
||
base = 0.1 + ord(chunk_id[-1]) % 7 / 100.0
|
||
return {
|
||
"id": chunk_id,
|
||
"embedding": [base] * DIM,
|
||
"product_id": product_id,
|
||
"product_name": "测试产品A",
|
||
"doc_type": doc_type,
|
||
"risk_level": "R2",
|
||
"source_doc_id": doc_id,
|
||
"source_version": version,
|
||
"effective_date": effective_date,
|
||
"chunk_text": text,
|
||
"chunk_no": int(chunk_id.split("_")[-1]),
|
||
}
|
||
|
||
|
||
@pytest.fixture(scope="module", autouse=True)
|
||
def seeded(client):
|
||
"""预置 5 条：3 条费率块（含 1 条未生效）+ 1 条规则块（他产品）+ 1 条风险块。"""
|
||
ms.insert_chunks(client, [
|
||
_row("PROD-TEST_0", "申购费率 1.5%，持有两年免赎回费"),
|
||
_row("PROD-TEST_1", "管理费每年 0.5%，托管费 0.1%"),
|
||
_row("PROD-TEST_2", "申购费率打一折", effective_date="2099-12-31"), # 未生效
|
||
_row("PROD-OTHER_0", "定投扣款日可修改", product_id="PROD-OTHER", doc_type="rule"),
|
||
_row("PROD-TEST_3", "本产品风险等级 R2，适合稳健型客户", doc_type="risk"),
|
||
])
|
||
|
||
|
||
class TestCollection:
|
||
def test_ensure_collection_idempotent(self, client):
|
||
ms.ensure_collection(client, dim=DIM) # 已存在 → 直接跳过不报错
|
||
assert client.has_collection(ms.COLLECTION_NAME)
|
||
|
||
|
||
class TestInsert:
|
||
def test_empty_rows_short_circuit(self, client):
|
||
assert ms.insert_chunks(client, []) == 0
|
||
|
||
def test_upsert_overwrites_same_id(self, client):
|
||
# 同 id 重写：upsert 覆盖而非报错/重复
|
||
row = _row("PROD-TEST_0", "申购费率 1.5%（改版后文案）")
|
||
ms.insert_chunks(client, [row])
|
||
hits = ms.search_kb(client, row["embedding"], top_k=1, product_id="PROD-TEST")
|
||
assert hits[0]["chunk_text"] == "申购费率 1.5%（改版后文案）"
|
||
|
||
|
||
class TestSearch:
|
||
def test_returns_source_refs_and_fields(self, client):
|
||
hits = ms.search_kb(client, [0.12] * DIM, top_k=2)
|
||
assert 1 <= len(hits) <= 2
|
||
for h in hits:
|
||
# 溯源字段必填（03-milvus-collections.md §2.3）
|
||
assert h["source_doc_id"] and h["source_version"]
|
||
assert h["product_id"] and h["chunk_text"]
|
||
assert isinstance(h["score"], float)
|
||
|
||
def test_effective_date_filter_excludes_future(self, client):
|
||
# §2.3：禁止跨 effective_date 过期文档——未生效块（2099）必须被排除
|
||
hits = ms.search_kb(client, [0.11] * DIM, top_k=10)
|
||
assert all(h["effective_date"] <= "2026-09-07" for h in hits)
|
||
|
||
def test_explicit_effective_on_date(self, client):
|
||
# effective_on 传历史日期：晚于该日的已生效块（2026-01-01 以外的）也排除
|
||
hits = ms.search_kb(client, [0.12] * DIM, top_k=10, effective_on="2020-01-01")
|
||
assert hits == []
|
||
|
||
def test_scalar_filter_product_id(self, client):
|
||
hits = ms.search_kb(client, [0.12] * DIM, top_k=10, product_id="PROD-OTHER")
|
||
assert len(hits) == 1
|
||
assert hits[0]["product_id"] == "PROD-OTHER"
|
||
|
||
def test_scalar_filter_doc_type(self, client):
|
||
hits = ms.search_kb(client, [0.12] * DIM, top_k=10, doc_type="risk")
|
||
assert len(hits) == 1
|
||
assert hits[0]["doc_type"] == "risk"
|