pydantic-ai-testing

Compare original and translation side by side

🇺🇸

Original

English

🇨🇳

Translation

Chinese

Testing PydanticAI Agents

测试PydanticAI Agents

TestModel (Deterministic Testing)

TestModel（确定性测试）

Use

TestModel

for tests without API calls:

python

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)

在无API调用的测试中使用

TestModel

：

python

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

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

    # 测试时用TestModel覆盖原模型
    result = agent.run_sync('Hello', model=TestModel())

    # TestModel会根据output_type生成确定性输出
    assert isinstance(result.output, str)

TestModel Configuration

TestModel 配置

python

from pydantic_ai.models.test import TestModel

python

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)

自定义结构化输出（适用于指定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'

from pydantic import BaseModel

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

Seed for reproducible random output

用于生成可复现随机输出的种子

model = TestModel(seed=42)

Force tool calls

强制调用工具

model = TestModel(call_tools=['my_tool', 'another_tool'])

undefined

model = TestModel(call_tools=['my_tool', 'another_tool'])

undefined

Override Context Manager

覆盖上下文管理器

python

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

python

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):
        # 所有运行都会使用TestModel和mock_deps
        result = agent.run_sync('Hello')
        assert result.output

FunctionModel (Custom Logic)

FunctionModel（自定义逻辑）

For complete control over model responses:

python

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

如需完全控制模型响应：

python

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:
    """自定义模型，可检查消息并返回响应。"""
    # 获取最后一条用户消息
    last_msg = messages[-1]

    # 返回自定义响应
    return ModelResponse(parts=[TextPart('Custom response')])

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

FunctionModel with Tool Calls

带工具调用的FunctionModel

python

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

python

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

def model_with_tools(
    messages: list[ModelMessage],
    info: AgentInfo
) -> ModelResponse:
    # 首次请求：调用工具
    if len(messages) == 1:
        return ModelResponse(parts=[
            ToolCallPart(
                tool_name='get_data',
                args='{"id": 123}'
            )
        ])

    # 工具响应后：返回最终结果
    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 Cassettes (Recorded API Calls)

VCR Cassettes（录制API调用）

Record and replay real LLM API interactions:

python

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

录制并重放真实LLM API交互：

python

import pytest

@pytest.mark.vcr
def test_with_recorded_response():
    """使用tests/cassettes/中的录制好的cassette"""
    agent = Agent('openai:gpt-4o')
    result = agent.run_sync('Hello')
    assert 'hello' in result.output.lower()

To record/update cassettes:

录制/更新cassettes：

uv run pytest --record-mode=rewrite tests/test_file.py


Cassette files are stored in `tests/cassettes/` as YAML.


Cassette文件以YAML格式存储在`tests/cassettes/`目录中。

Inline Snapshots

内联快照

Assert expected outputs with auto-updating snapshots:

python

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

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

python

from inline_snapshot import snapshot

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

    # 首次运行：创建快照
    # 后续运行：与快照进行断言
    assert result.output == snapshot('expected output here')

Update snapshots:

更新快照：

uv run pytest --inline-snapshot=fix

undefined

undefined

Testing Tools

测试工具

python

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

python

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}"

    # 强制TestModel调用该工具
    result = agent.run_sync(
        'Use my_tool',
        model=TestModel(call_tools=['my_tool'])
    )

    assert tool_called

Testing with Dependencies

依赖项测试

python

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

python

from dataclasses import dataclass
from unittest.mock import AsyncMock

@dataclass
class Deps:
    api: ApiClient

def test_tool_with_deps():
    # 创建模拟依赖项
    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()

Capture Messages

捕获消息

Inspect all messages in a run:

python

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

检查运行过程中的所有消息：

python

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)

undefined

for msg in messages: print(msg)

undefined

Testing Patterns Summary

测试模式汇总

Scenario	Approach
Unit tests without API	`TestModel()`
Custom model logic	`FunctionModel(func)`
Recorded real responses	`@pytest.mark.vcr`
Assert output structure	`inline_snapshot`
Test tools are called	`TestModel(call_tools=[...])`
Mock dependencies	`agent.override(deps=...)`

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

pytest Configuration

pytest 配置

Typical

pyproject.toml

toml

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

Run tests:

bash

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

典型的

pyproject.toml

配置：

toml

[tool.pytest.ini_options]
testpaths = ["tests"]
asyncio_mode = "auto"  # 用于异步测试

运行测试：

bash

uv run pytest tests/test_agent.py -v
uv run pytest --inline-snapshot=fix  # 更新快照