Pydantic-AI项目单元测试最佳实践指南-优快云博客

本文链接：https://blog.youkuaiyun.com/gitblog_00133/article/details/148417034

Pydantic-AI项目单元测试最佳实践指南

pydantic-ai Agent Framework / shim to use Pydantic with LLMs 项目地址: https://gitcode.com/gh_mirrors/py/pydantic-ai

引言

在开发基于Pydantic-AI的项目时，编写有效的单元测试是确保代码质量的关键环节。本文将深入探讨如何为Pydantic-AI应用编写高质量的单元测试，涵盖测试策略选择、工具使用以及实际案例演示。

测试策略概述

Pydantic-AI应用的测试与其他Python项目类似，但有一些特殊考虑：

测试框架选择：推荐使用pytest作为主要测试框架
断言优化：对于复杂断言，可以使用inline-snapshot和dirty-equals等工具简化
模型隔离：使用测试专用模型避免真实LLM调用

核心测试工具

TestModel的使用

TestModel是Pydantic-AI提供的测试专用模型，它能够：

自动调用代理中的所有工具
根据返回类型生成纯文本或结构化响应
基于注册工具的JSON Schema生成有效数据

from pydantic_ai.models.test import TestModel

# 在测试中使用TestModel
with weather_agent.override(model=TestModel()):
    # 测试代码

FunctionModel的进阶使用

当需要更精细控制测试行为时，可以使用FunctionModel：

from pydantic_ai.models.function import FunctionModel

def custom_model_behavior(messages, info):
    # 自定义模型响应逻辑
    return ModelResponse(...)

# 在测试中使用
with weather_agent.override(model=FunctionModel(custom_model_behavior)):
    # 测试代码

实战测试案例

天气应用测试示例

假设我们有一个天气预测应用，以下是测试要点：

测试准备：

pytestmark = pytest.mark.anyio
models.ALLOW_MODEL_REQUESTS = False  # 禁止真实模型请求

基本测试：

async def test_forecast():
    with capture_run_messages() as messages:
        with weather_agent.override(model=TestModel()):
            await run_weather_forecast(test_data, test_conn)
    
    # 验证数据库存储结果
    assert forecast == expected_result
    
    # 验证消息交互流程
    assert messages == expected_message_sequence

自定义行为测试：

def call_weather_forecast(messages, info):
    # 解析用户输入中的日期
    date = extract_date_from_prompt(messages)
    return ModelResponse(
        parts=[ToolCallPart('weather_forecast', 
              {'location': 'London', 'forecast_date': date})]
    )

async def test_custom_behavior():
    with weather_agent.override(model=FunctionModel(call_weather_forecast)):
        await run_weather_forecast(custom_prompt, test_conn)

高级测试技巧

使用pytest fixture简化测试

对于需要重复设置测试模型的场景，可以使用fixture：

@pytest.fixture
def test_model_fixture():
    with weather_agent.override(model=TestModel()):
        yield

async def test_with_fixture(test_model_fixture):
    # 测试代码无需重复设置模型
    result = await weather_agent.run(test_prompt)
    assert validate_result(result)

消息捕获与验证

使用capture_run_messages可以捕获和验证模型交互过程：

with capture_run_messages() as messages:
    # 执行测试代码
    await agent.run(prompt)

# 验证消息序列
assert len(messages) == expected_message_count
assert messages[0].parts[0].content == expected_system_prompt