Files
group_xinghuo_jinrong/tests/test_milvus_service.py
T

116 lines
4.6 KiB
Python
Raw Normal View History

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