group_xinghuo_jinrong/tests/test_rag_service.py

"""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("任意")