TestModelとの併用

from pydantic_ai import Agent
from pydantic_ai.models.test import TestModel
 
agent = Agent(
    'google-gla:gemini-3-flash-preview',
    instructions='短い日本語で答えてください。',
)
 
def test_basic_run():
    # テスト時だけ TestModel に差し替え
    with agent.override(model=TestModel(custom_output_text='テスト用の固定応答')):
        result = agent.run_sync('何でも良いので答えて')
    assert result.output == 'テスト用の固定応答'

from dataclasses import dataclass
from pydantic_ai import Agent, RunContext
from pydantic_ai.models.test import TestModel
 
@dataclass
class CityDb:
    cities: dict[str, str]
 
agent = Agent(
    'google-gla:gemini-3-flash-preview',
    deps_type=CityDb,
    instructions='ユーザーの国名から首都を答えてください。',
)
 
@agent.instructions
def known_cities(ctx: RunContext[CityDb]) -> str:
    return f'参照可能な国: {", ".join(ctx.deps.cities.keys())}'
 
def test_instructions_includes_country_list():
    fake_deps = CityDb(cities={'X国': 'X市', 'Y国': 'Y市'})
    with agent.override(model=TestModel(), deps=fake_deps):
        result = agent.run_sync('テスト', deps=fake_deps)
    # TestModel が組み立てた "ツール要約" 文字列を見て、
    # known_cities が呼ばれた痕跡 (LLM への system message に X国/Y国 が出ている) を確認
    # → メッセージ履歴を見る方が確実
    msgs = result.all_messages()
    sys_text = ''.join(
        part.content for m in msgs for part in m.parts if hasattr(part, 'content')
    )
    assert 'X国' in sys_text and 'Y国' in sys_text

from pydantic_ai import Agent, ModelRetry
from pydantic_ai.exceptions import UnexpectedModelBehavior
from pydantic_ai.models.test import TestModel
 
agent = Agent('google-gla:gemini-3-flash-preview', instructions='短く答えて。')
 
@agent.output_validator
def no_secrets(output: str) -> str:
    if '秘密' in output:
        raise ModelRetry('「秘密」を含めずに言い換えて')
    return output
 
def test_validator_rejects_then_passes():
    # 1 回目は NG、2 回目は OK を返す TestModel
    test_model = TestModel(custom_output_text='秘密の手順')   # NG な文字列
    with agent.override(model=test_model):
        # retries のデフォルトは 1。TestModel は毎回同じ文字列を返すので、
        # 2 回目も同じ NG → 上限を超えて UnexpectedModelBehavior が飛ぶことを検証
        try:
            agent.run_sync('教えて')
            assert False, '例外が飛ぶはず'
        except UnexpectedModelBehavior as e:
            assert '秘密' in str(e) or 'retry' in str(e).lower()

from pydantic_ai.messages import ModelResponse, TextPart
from pydantic_ai.models.function import FunctionModel, AgentInfo
 
calls = {'n': 0}
 
def model_logic(messages, info: AgentInfo) -> ModelResponse:
    calls['n'] += 1
    # 1 回目は NG、2 回目以降は OK を返す
    if calls['n'] == 1:
        return ModelResponse(parts=[TextPart(content='秘密の手順を説明します')])
    return ModelResponse(parts=[TextPart(content='ログイン手順を一般的に説明します')])
 
def test_validator_passes_after_retry():
    calls['n'] = 0
    with agent.override(model=FunctionModel(model_logic)):
        result = agent.run_sync('教えて')
    assert '秘密' not in result.output
    assert calls['n'] == 2  # 1 回目で ModelRetry → 2 回目で OK

import pytest
from dataclasses import dataclass
 
@dataclass
class FakeDeps:
    db: dict[str, str]
 
@pytest.fixture
def fake_deps() -> FakeDeps:
    return FakeDeps(db={'Tokyo': 'sunny', 'Paris': 'rainy'})
 
@pytest.fixture
def fake_agent():
    # 全テストで共通の TestModel
    return TestModel(custom_output_text='ok')
 
def test_run_with_fake(fake_deps, fake_agent):
    with agent.override(model=fake_agent, deps=fake_deps):
        result = agent.run_sync('Tokyo の天気は?', deps=fake_deps)
    assert result.output == 'ok'