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 get_query_log_by_trace(self, trace_id):
        return {
            "staff_id": "STAFF-A",
            "generated_sql": "SELECT COUNT(*) AS c FROM core_customer",
            "exec_status": "success",
        }


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_trade_flow_ambiguity(self):
        agent = AnalystAgent(llm=FakeLLM("SELECT 1", []), repo=FakeRepo())
        resp = agent.run("近30天交易流水是多少", ctx(["analyst"]))
        self.assertEqual(resp.status, "clarify")
        self.assertIn("流水", resp.answer)

    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_query_only(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(resp.answer, "")
        self.assertEqual(len(repo.logged), 1)
        self.assertEqual(agent.llm.calls, 1)

    def test_success_with_interpret_flag(self):
        repo = FakeRepo(rows=[[33]], columns=["c"])
        llm = FakeLLM("SELECT COUNT(*) AS c FROM core_customer", ["共 33 个客户"])
        agent = AnalystAgent(llm=llm, repo=repo)
        resp = agent.run("客户总数是多少", ctx(["analyst"]), interpret=True)
        self.assertEqual(resp.status, "success")
        self.assertIn("33", resp.answer)
        self.assertEqual(llm.calls, 2)

    def test_interpret_from_snapshot(self):
        from app.model.analyst_schemas import InterpretRequest, Meta, TableData

        class AnswerOnlyLLM:
            def complete(self, messages, temperature=0, max_tokens=2048):
                return "共 33 个客户", {"prompt_tokens": 1, "completion_tokens": 1}

        repo = FakeRepo(rows=[[33]], columns=["c"])
        agent = AnalystAgent(llm=AnswerOnlyLLM(), repo=repo)
        req = InterpretRequest(
            question="客户总数是多少",
            status="success",
            table=TableData(columns=["c"], rows=[[33]]),
            sql="SELECT COUNT(*) AS c FROM core_customer",
            meta=Meta(row_count=1),
        )
        resp = agent.interpret(req, ctx(["analyst"]))
        self.assertEqual(resp.status, "success")
        self.assertIn("33", resp.answer)

    def test_interpret_echo_deny_without_llm(self):
        from app.model.analyst_schemas import InterpretRequest

        llm = FakeLLM("SELECT 1", [])
        agent = AnalystAgent(llm=llm, repo=FakeRepo())
        req = InterpretRequest(
            question="查别人",
            status="deny",
            answer="无法执行：权限不足",
        )
        resp = agent.interpret(req, ctx(["customer"], "CUST-1", customer_id="CUST-1"))
        self.assertEqual(resp.status, "deny")
        self.assertEqual(resp.answer, "无法执行：权限不足")
        self.assertEqual(llm.calls, 0)

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

    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"]), interpret=True)
        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"),
            interpret=True,
        )
        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, 0)

    def test_sql_exec_failure_escalates(self):
        class FailRepo(FakeRepo):
            def execute_readonly(self, sql):
                raise TimeoutError("read timed out")

        class NoCache:
            def permission_fingerprint(self, *a, **k):
                return "fp"

            def sql_hash(self, sql):
                return "hash"

            def get_result(self, *a, **k):
                return None

            def set_result(self, *a, **k):
                pass

        agent = AnalystAgent(
            llm=FakeLLM("SELECT COUNT(*) AS c FROM core_customer", []),
            repo=FailRepo(),
            templates=TemplateService(templates=[]),
            cache=NoCache(),
        )
        resp = agent.run("客户总数", ctx(["analyst"]))
        self.assertEqual(resp.status, "escalate")
        self.assertIn(resp.error_code, ("EXEC_TIMEOUT", "EXEC_ERROR"))

    def test_sample_by_trace(self):
        repo = FakeRepo(rows=[[1]], columns=["c"])
        agent = AnalystAgent(llm=FakeLLM("SELECT 1", []), repo=repo)
        resp = agent.sample_by_trace("trace-1", ctx(["analyst"], subject="STAFF-A"))
        self.assertEqual(resp.status, "success")
        self.assertIn("LIMIT", resp.sql)

    def test_escalate_audit(self):
        repo = FakeRepo()
        agent = AnalystAgent(llm=FakeLLM("SELECT 1", []), repo=repo)
        out = agent.escalate("trace-x", "q", "timeout", ctx(["analyst"]))
        self.assertTrue(out["ok"])


@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)