group_xinghuo_jinrong/tests/test_wave6_analyst_agent.py

"""analyst_agent 编排测试（Wave 6）。"""
import unittest

import pytest

from app.api.analyst_auth_adapter import AnalystAuthContext
from app.service.analyst_agent import AnalystAgent
from app.service.template_service import QueryTemplate, TemplateService


class FakeLLM:
    def __init__(self, sql, answers):
        self.sql = sql
        self.answers = list(answers)
        self.calls = 0

    def complete(self, messages, temperature=0, max_tokens=2048):
        self.calls += 1
        usage = {"prompt_tokens": 10, "completion_tokens": 10}
        if self.calls == 1:
            return self.sql, usage
        ans = self.answers.pop(0) if self.answers else "无解读"
        return ans, usage


class FakeRepo:
    def __init__(self, rows=(), columns=(), scope=None):
        self.rows = list(rows)
        self.columns = list(columns)
        self.scope = scope or []
        self.logged = []

    def resolve_advisor_scope(self, sid):
        return self.scope

    def execute_readonly(self, sql):
        return {"columns": self.columns, "rows": self.rows}

    def get_data_as_of(self):
        return "2026-09-04"

    def log_query(self, **kw):
        self.logged.append(kw)

    def log_audit(self, **kw):
        pass

    def list_published_templates(self):
        return []


def ctx(roles, subject="STAFF-A", *, token_type="staff", customer_id=None):
    return AnalystAuthContext(
        subject_id=subject,
        token_type=token_type,
        roles=roles,
        customer_id=customer_id,
    )


class TestAgentOrchestration(unittest.TestCase):
    def test_clarify_ambiguity(self):
        agent = AnalystAgent(llm=FakeLLM("SELECT 1", []), repo=FakeRepo())
        resp = agent.run("我名下的规模是多少", ctx(["analyst"]))
        self.assertEqual(resp.status, "clarify")

    def test_success(self):
        repo = FakeRepo(rows=[[33]], columns=["c"])
        agent = AnalystAgent(
            llm=FakeLLM("SELECT COUNT(*) AS c FROM core_customer", ["共 33 个客户"]),
            repo=repo,
        )
        resp = agent.run("客户总数是多少", ctx(["analyst"]))
        self.assertEqual(resp.status, "success")
        self.assertEqual(resp.table.rows, [[33]])
        self.assertEqual(len(repo.logged), 1)

    def test_deny_bad_sql(self):
        agent = AnalystAgent(llm=FakeLLM("INSERT INTO core_customer VALUES (1)", []), repo=FakeRepo())
        resp = agent.run("删库", ctx(["analyst"]))
        self.assertEqual(resp.status, "deny")
        self.assertEqual(resp.error_code, "SQL_NOT_SELECT")

    def test_degrade_wrong_number(self):
        repo = FakeRepo(rows=[[33]], columns=["c"])
        agent = AnalystAgent(
            llm=FakeLLM("SELECT COUNT(*) FROM core_customer", ["共 999 个客户", "共 999 个客户"]),
            repo=repo,
        )
        resp = agent.run("客户总数", ctx(["analyst"]))
        self.assertEqual(resp.status, "degrade")

    def test_advisor_out_of_scope_deny(self):
        repo = FakeRepo(scope=["CUST-1001"])
        agent = AnalystAgent(
            llm=FakeLLM("SELECT * FROM core_holding WHERE customer_id='CUST-1004'", []),
            repo=repo,
        )
        resp = agent.run("查 CUST-1004 持仓", ctx(["advisor"], "STAFF-B"))
        self.assertEqual(resp.status, "deny")
        self.assertEqual(resp.error_code, "AUTH_403_NOT_ASSIGNED")

    def test_customer_self_success(self):
        repo = FakeRepo(rows=[[2]], columns=["cnt"])
        agent = AnalystAgent(
            llm=FakeLLM(
                "SELECT COUNT(*) AS cnt FROM core_trade WHERE customer_id='CUST-9527'",
                ["近阶段共有 2 笔交易"],
            ),
            repo=repo,
        )
        resp = agent.run(
            "我有多少笔交易",
            ctx(["customer"], "CUST-9527", token_type="customer", customer_id="CUST-9527"),
        )
        self.assertEqual(resp.status, "success")
        self.assertIn("AI 分析有风险", resp.answer)

    def test_template_hit_skips_llm_sql(self):
        tpl = QueryTemplate(
            template_key="customer_total_count",
            template_sql="SELECT COUNT(*) AS c FROM core_customer",
            params_schema={"match_all": ["客户", "总数"]},
        )
        repo = FakeRepo(rows=[[33]], columns=["c"])
        llm = FakeLLM("SELECT 1", ["共 33 个客户"])
        agent = AnalystAgent(
            llm=llm,
            repo=repo,
            templates=TemplateService(templates=[tpl]),
        )
        resp = agent.run("客户总数是多少", ctx(["analyst"]))
        self.assertEqual(resp.status, "success")
        self.assertTrue(resp.meta.template_hit)
        self.assertEqual(resp.meta.template_key, "customer_total_count")
        self.assertIn("COUNT(*)", resp.sql)
        self.assertEqual(llm.calls, 1)


@pytest.mark.integration
class TestAgentReal(unittest.TestCase):
    @pytest.mark.skip(reason="需要真实 MySQL + DeepSeek Key")
    def test_real_end_to_end(self):
        from app.service.analytics_repo import AnalyticsRepo
        from app.service.llm import DeepSeekLLM

        agent = AnalystAgent(llm=DeepSeekLLM(), repo=AnalyticsRepo())
        resp = agent.run("客户总数是多少", ctx(["analyst"]))
        self.assertIn(resp.status, ("success", "degrade"))
        self.assertTrue(resp.sql)
        self.assertGreater(resp.meta.row_count, 0)