pydantic-ai-testing

测试 PydanticAI Agent

TestModel（确定性测试）

使用 TestModel 进行无需调用 API 的测试：

import pytest
from pydantic_ai import Agent
from pydantic_ai.models.test import TestModel

def test_agent_basic():
    agent = Agent('openai:gpt-4o')

    # Override with TestModel for testing
    result = agent.run_sync('Hello', model=TestModel())

    # TestModel generates deterministic output based on output_type
    assert isinstance(result.output, str)

TestModel 配置

from pydantic_ai.models.test import TestModel

# Custom text output
model = TestModel(custom_output_text='Custom response')
result = agent.run_sync('Hello', model=model)
assert result.output == 'Custom response'

# Custom structured output (for output_type agents)
from pydantic import BaseModel

class Response(BaseModel):
    message: str
    score: int

agent = Agent('openai:gpt-4o', output_type=Response)
model = TestModel(custom_output_args={'message': 'Test', 'score': 42})
result = agent.run_sync('Hello', model=model)
assert result.output.message == 'Test'

# Seed for reproducible random output
model = TestModel(seed=42)

# Force tool calls
model = TestModel(call_tools=['my_tool', 'another_tool'])

Override 上下文管理器

from pydantic_ai import Agent
from pydantic_ai.models.test import TestModel

agent = Agent('openai:gpt-4o', deps_type=MyDeps)

def test_with_override():
    mock_deps = MyDeps(db=MockDB())

    with agent.override(model=TestModel(), deps=mock_deps):
        # All runs use TestModel and mock_deps
        result = agent.run_sync('Hello')
        assert result.output

FunctionModel（自定义逻辑）

完全控制模型响应：

from pydantic_ai import Agent, ModelMessage, ModelResponse, TextPart
from pydantic_ai.models.function import AgentInfo, FunctionModel

def custom_model(
    messages: list[ModelMessage],
    info: AgentInfo
) -> ModelResponse:
    """Custom model that inspects messages and returns response."""
    # Access the last user message
    last_msg = messages[-1]

    # Return custom response
    return ModelResponse(parts=[TextPart('Custom response')])

agent = Agent(FunctionModel(custom_model))
result = agent.run_sync('Hello')

带工具调用的 FunctionModel

from pydantic_ai import ToolCallPart, ModelResponse
from pydantic_ai.models.function import AgentInfo, FunctionModel

def model_with_tools(
    messages: list[ModelMessage],
    info: AgentInfo
) -> ModelResponse:
    # First request: call a tool
    if len(messages) == 1:
        return ModelResponse(parts=[
            ToolCallPart(
                tool_name='get_data',
                args='{"id": 123}'
            )
        ])

    # After tool response: return final result
    return ModelResponse(parts=[TextPart('Done with tool result')])

agent = Agent(FunctionModel(model_with_tools))

@agent.tool_plain
def get_data(id: int) -> str:
    return f"Data for {id}"

result = agent.run_sync('Get data')

VCR 录制带（已记录的 API 调用）

录制并重放真实的大模型 API 交互：

import pytest

@pytest.mark.vcr
def test_with_recorded_response():
    """Uses recorded cassette from tests/cassettes/"""
    agent = Agent('openai:gpt-4o')
    result = agent.run_sync('Hello')
    assert 'hello' in result.output.lower()

# To record/update cassettes:
# uv run pytest --record-mode=rewrite tests/test_file.py

录制带文件以 YAML 格式存储在 tests/cassettes/ 中。

内联快照

使用自动更新的快照断言预期输出：

from inline_snapshot import snapshot

def test_agent_output():
    result = agent.run_sync('Hello', model=TestModel())

    # First run: creates snapshot
    # Subsequent runs: asserts against it
    assert result.output == snapshot('expected output here')

# Update snapshots:
# uv run pytest --inline-snapshot=fix

门禁：VCR 录制带和内联快照

录制或修复会重写磁盘上的文件。请按以下顺序操作；不要跳过步骤。

1. 重放阶段（不带录制/修复标志）： 对目标路径运行 uv run pytest；确保全部通过（或已了解失败原因，且与要刷新的产物无关）。
2. 锁定范围： 确认要更新的 tests/cassettes/ 下的录制带或 snapshot(...) 断言；确认只有这些文件应发生更改。
3. 录制或修复： 仅对该路径运行一条范围限定的命令：uv run pytest --record-mode=rewrite … 或 uv run pytest --inline-snapshot=fix …。
4. 后置条件： 再次运行相同的测试，不带录制/修复标志；确保全部通过。检查 git diff ——只应有预期的 .yaml / 快照更改。

如果第 4 步失败，请还原意外的差异，并在重新录制前修复测试或模型。

测试工具

from pydantic_ai import Agent, RunContext
from pydantic_ai.models.test import TestModel

def test_tool_is_called():
    agent = Agent('openai:gpt-4o')
    tool_called = False

    @agent.tool_plain
    def my_tool(x: int) -> str:
        nonlocal tool_called
        tool_called = True
        return f"Result: {x}"

    # Force TestModel to call the tool
    result = agent.run_sync(
        'Use my_tool',
        model=TestModel(call_tools=['my_tool'])
    )

    assert tool_called

使用依赖项进行测试

from dataclasses import dataclass
from unittest.mock import AsyncMock

@dataclass
class Deps:
    api: ApiClient

def test_tool_with_deps():
    # Create mock dependency
    mock_api = AsyncMock()
    mock_api.fetch.return_value = {'data': 'test'}

    agent = Agent('openai:gpt-4o', deps_type=Deps)

    @agent.tool
    async def fetch_data(ctx: RunContext[Deps]) -> dict:
        return await ctx.deps.api.fetch()

    with agent.override(
        model=TestModel(call_tools=['fetch_data']),
        deps=Deps(api=mock_api)
    ):
        result = agent.run_sync('Fetch data')

    mock_api.fetch.assert_called_once()

捕获消息

检查一次运行中的所有消息：

from pydantic_ai import Agent, capture_run_messages

agent = Agent('openai:gpt-4o')

with capture_run_messages() as messages:
    result = agent.run_sync('Hello', model=TestModel())

# Inspect captured messages
for msg in messages:
    print(msg)

测试模式总结

场景	方法
无需 API 的单元测试	TestModel()
自定义模型逻辑	FunctionModel(func)
录制真实响应	@pytest.mark.vcr
断言输出结构	inline_snapshot
测试是否调用工具	TestModel(call_tools=[...])
模拟依赖项	agent.override(deps=...)

pytest 配置

典型的 pyproject.toml：

[tool.pytest.ini_options]
testpaths = ["tests"]
asyncio_mode = "auto"  # For async tests

运行测试：

uv run pytest tests/test_agent.py -v
uv run pytest --inline-snapshot=fix  # Update snapshots