Agent 测试金字塔（Testing Pyramid for AI Agents）

# 基于 pytest + unittest.mock 演示（截至 2026-03）

from typing import List, Dict, Optional
from dataclasses import dataclass
from unittest.mock import Mock
import json

# ========== Agent 最小实现 ==========
@dataclass
class ToolCall:
    """工具调用结构体"""
    name: str
    arguments: Dict

class SimpleAgent:
    """最小 Agent：解析工具调用 → 执行 → 返回结果"""
    def __init__(self, model, tools: Dict, max_iterations: int = 10):
        self.model = model
        self.tools = tools
        self.max_iterations = max_iterations

    def run(self, question: str) -> str:
        history = [{"role": "user", "content": question}]
        for _ in range(self.max_iterations):
            response = self.model.generate(history)
            tool_call = self._parse_tool_call(response)
            if tool_call is None:
                return response  # 无工具调用，直接返回答案
            if tool_call.name not in self.tools:
                return f"错误：工具 '{tool_call.name}' 不存在"
            result = self.tools[tool_call.name](tool_call.arguments)
            history.append({"role": "assistant", "content": response})
            history.append({"role": "tool", "content": str(result)})
        return "已达最大迭代次数"

    def _parse_tool_call(self, response: str) -> Optional[ToolCall]:
        if "TOOL:" not in response:
            return None
        try:
            parts = response.split("TOOL:")[1].split("ARGS:")
            name = parts[0].strip()
            args = json.loads(parts[1].strip()) if len(parts) > 1 else {}
            return ToolCall(name=name, arguments=args)
        except Exception:
            return None


# ========== 第一层：单元测试 ==========
# 特征：Mock LLM + Mock 工具，验证框架逻辑

class MockLLM:
    """返回预设响应的 Mock 模型"""
    def __init__(self, responses: List[str]):
        self.responses = responses
        self._idx = 0
    def generate(self, history) -> str:
        resp = self.responses[self._idx % len(self.responses)]
        self._idx += 1
        return resp

def test_max_iteration_limit():
    """Agent 必须在达到上限时停止，不能无限循环"""
    model = MockLLM(['TOOL: search ARGS: {"q": "test"}'] * 5)
    tool = Mock(return_value="结果")
    agent = SimpleAgent(model, {"search": tool}, max_iterations=3)
    result = agent.run("测试问题")
    assert result == "已达最大迭代次数"
    assert tool.call_count == 3  # 恰好调用 3 次就停止

def test_unknown_tool_handling():
    """调用不存在的工具时应返回错误信息"""
    model = MockLLM(['TOOL: nonexistent ARGS: {}'])
    agent = SimpleAgent(model, {}, max_iterations=5)
    result = agent.run("测试")
    assert "不存在" in result


# ========== 第二层：集成测试 ==========
# 特征：Mock LLM + 真实工具逻辑，验证工具协作

def real_database_query(args: Dict) -> str:
    """模拟真实数据库查询（实际项目中连接测试数据库）"""
    if "user" in args.get("sql", "").lower():
        return json.dumps({"id": 1, "name": "张三", "email": "zhang@example"})
    return json.dumps({})

def test_agent_database_interaction():
    """Agent 调用数据库工具后能正确提取信息"""
    model = MockLLM([
        'TOOL: db ARGS: {"sql": "SELECT * FROM users WHERE id=1"}',
        '用户张三的邮箱是 zhang@example'
    ])
    agent = SimpleAgent(model, {"db": real_database_query}, max_iterations=5)
    result = agent.run("查询用户1的信息")
    assert "张三" in result
    assert "zhang@example" in result


# ========== 第三层：评估驱动测试（伪代码示意） ==========
# 特征：真实 LLM + 真实工具，概率性判定

def evaluate_with_scoring(agent, test_cases, num_runs=10):
    """
    对每个测试用例运行 N 次，统计成功率和平均质量评分。
    实际项目中 judge_model 调用真实 LLM API。
    """
    results = {}
    for case in test_cases:
        scores = []
        for _ in range(num_runs):
            answer = agent.run(case["question"])
            # 实际中用 LLM-as-a-Judge 或规则检查打分
            score = 1.0 if case["keyword"] in answer else 0.0
            scores.append(score)
        success_rate = sum(s > 0.7 for s in scores) / num_runs
        avg_score = sum(scores) / len(scores)
        results[case["question"]] = {
            "success_rate": success_rate,
            "avg_score": avg_score
        }
    return results

# 评估驱动测试的判定标准：成功率 > 80% 且平均分 > 0.75
# 不是 assert == 某个值，而是 assert >= 某个阈值