126 lines
4.5 KiB
Python
126 lines
4.5 KiB
Python
"""T21-4 rag_service 单测：检索编排 + source_refs 溯源（monkeypatch 注入，不连外部）。
|
||
|
||
embedding 与 Milvus client 均经模块级 _embed/_client 注入点 mock；
|
||
Milvus 检索行为用假 client（FakeClient）模拟，search_kb 参数透传断言。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import pytest
|
||
|
||
from app.service import rag_service as rs
|
||
from app.service.embedding import EmbeddingError
|
||
|
||
|
||
class FakeClient:
|
||
"""假 Milvus client：记录 ensure/close 调用，search 透传给 stub。"""
|
||
|
||
def __init__(self, search_stub):
|
||
self.search_stub = search_stub
|
||
self.closed = False
|
||
self.last_kwargs: dict = {}
|
||
|
||
def close(self):
|
||
self.closed = True
|
||
|
||
|
||
def _client_with(monkeypatch, stub) -> FakeClient:
|
||
"""构造假 client 并注入 rs._client（embed 走假 8 维向量）。"""
|
||
client = FakeClient(stub)
|
||
monkeypatch.setattr(rs, "_client", lambda: client)
|
||
monkeypatch.setattr(rs, "_embed", lambda q: [0.1] * 8)
|
||
monkeypatch.setattr(
|
||
rs.milvus_service, "ensure_collection", lambda c, dim=None: None
|
||
)
|
||
monkeypatch.setattr(
|
||
rs.milvus_service,
|
||
"search_kb",
|
||
lambda c, vector, **kwargs: (client.last_kwargs.update(kwargs), stub(vector, **kwargs))[1],
|
||
)
|
||
return client
|
||
|
||
|
||
def _hit(idx: str, doc: str = "KB-P1", ver: str = "v1", pid: str = "PROD-1") -> dict:
|
||
return {
|
||
"id": idx,
|
||
"score": 0.9,
|
||
"product_id": pid,
|
||
"product_name": "产品一",
|
||
"doc_type": "fee",
|
||
"risk_level": "R2",
|
||
"source_doc_id": doc,
|
||
"source_version": ver,
|
||
"effective_date": "2026-09-01",
|
||
"chunk_text": "申购费率 1.5%",
|
||
"chunk_no": 0,
|
||
}
|
||
|
||
|
||
class TestSearch:
|
||
def test_basic_flow_and_source_refs(self, monkeypatch):
|
||
_client_with(monkeypatch, lambda v, **k: [_hit("P1_0"), _hit("P1_1", ver="v2")])
|
||
out = rs.search_knowledge("基金申购费率")
|
||
assert len(out["results"]) == 2
|
||
# 同文档不同版本 → 溯源按 (doc, ver, pid) 去重，应为 2 条
|
||
assert len(out["source_refs"]) == 2
|
||
assert out["source_refs"][0]["source_doc_id"] == "KB-P1"
|
||
|
||
def test_blank_query_short_circuits(self, monkeypatch):
|
||
# 空白 query：不发 embedding、不连 Milvus
|
||
def bad_client():
|
||
raise AssertionError("空白 query 不应触发 Milvus 连接")
|
||
|
||
monkeypatch.setattr(rs, "_client", bad_client)
|
||
monkeypatch.setattr(rs, "_embed", lambda q: [] if False else (_ for _ in ()).throw(AssertionError("不应向量化")))
|
||
out = rs.search_knowledge(" ")
|
||
assert out == {"query": "", "results": [], "source_refs": []}
|
||
|
||
def test_filter_params_passthrough(self, monkeypatch):
|
||
client = _client_with(monkeypatch, lambda v, **k: [])
|
||
rs.search_knowledge("定投规则", product_id="PROD-9", doc_type="rule", top_k=5)
|
||
assert client.last_kwargs["product_id"] == "PROD-9"
|
||
assert client.last_kwargs["doc_type"] == "rule"
|
||
assert client.last_kwargs["top_k"] == 5
|
||
|
||
def test_query_whitespace_stripped(self, monkeypatch):
|
||
_client_with(monkeypatch, lambda v, **k: [])
|
||
out = rs.search_knowledge(" 赎回到账时间 ")
|
||
assert out["query"] == "赎回到账时间"
|
||
|
||
|
||
class TestSourceRefs:
|
||
def test_dedup_same_doc_version(self, monkeypatch):
|
||
_client_with(monkeypatch, lambda v, **k: [_hit("P1_0"), _hit("P1_1")])
|
||
out = rs.search_knowledge("费率")
|
||
assert len(out["source_refs"]) == 1
|
||
|
||
def test_multi_product_refs(self, monkeypatch):
|
||
_client_with(
|
||
monkeypatch,
|
||
lambda v, **k: [_hit("P1_0"), _hit("P2_0", doc="KB-P2", pid="PROD-2")],
|
||
)
|
||
out = rs.search_knowledge("费率")
|
||
assert len(out["source_refs"]) == 2
|
||
pids = {r["product_id"] for r in out["source_refs"]}
|
||
assert pids == {"PROD-1", "PROD-2"}
|
||
|
||
|
||
class TestFailurePropagation:
|
||
"""失败口径：异常原样上抛，不吞不降级（防止 LLM 编造回答）。"""
|
||
|
||
def test_embed_error_propagates(self, monkeypatch):
|
||
def boom(q):
|
||
raise EmbeddingError("Ollama 连接失败")
|
||
|
||
monkeypatch.setattr(rs, "_embed", boom)
|
||
with pytest.raises(EmbeddingError):
|
||
rs.search_knowledge("任意")
|
||
|
||
def test_milvus_error_propagates(self, monkeypatch):
|
||
def stub(vector, **kwargs):
|
||
raise RuntimeError("milvus down")
|
||
|
||
_client_with(monkeypatch, stub)
|
||
with pytest.raises(RuntimeError):
|
||
rs.search_knowledge("任意")
|