Pydantic AI 測試與 Mocking (單元測試)

在傳統的軟體開發中,我們習慣撰寫單元測試 (Unit Tests) 來確保程式碼的正確性。然而,測試 AI Agent 往往會遇到兩個挑戰:

  1. 不確定性:LLM 每次產生的文字可能都不一樣。
  2. 成本與速度:每次跑測試如果都要真實呼叫 OpenAI 或 Gemini API,不僅測試執行速度極慢,還會消耗大量的 API 額度 (Tokens)。

為了解決這些問題,Pydantic AI 在設計之初就將「可測試性 (Testability)」放在核心位置。它提供了極佳的 Mocking 機制,讓你可以快速、免費且確定性地測試你的 Agent 邏輯。

覆寫模型:使用 TestModel

Pydantic AI 提供了一個專為測試設計的虛擬模型 TestModel

當你在測試環境中使用 TestModel 來替換原本的真實模型時,你可以強制作弊,預先設定好模型會回傳什麼內容。這對於測試 Agent 是否有正確呼叫工具、或是後端程式碼是否能正確解析 Agent 的輸出非常有幫助。

測試範例

假設我們有一個查詢匯率的 Agent,它具備一個查詢即時匯率的工具函數:

# --- 產品程式碼 (agent.py) ---
from pydantic_ai import Agent

currency_agent = Agent('openai:gpt-4o-mini')

@currency_agent.tool
def get_exchange_rate(currency_code: str) -> float:
    # 這裡原本可能是真實呼叫外部 API
    return 32.5

# --- 測試程式碼 (test_agent.py) ---
from pydantic_ai.models.test import TestModel
# 匯入我們寫好的 Agent
# from agent import currency_agent 

def test_currency_agent_tool_calling():
    # 建立一個測試用的虛擬模型
    test_model = TestModel()
    
    # 告訴虛擬模型:當你被執行時,請強迫自己去呼叫 get_exchange_rate 工具,
    # 並且使用 'USD' 作為工具的參數
    test_model.custom_result_args = {'get_exchange_rate': {'currency_code': 'USD'}}
    
    # 使用 override 裝飾器或上下文管理器 (Context Manager)
    # 在執行這個區塊時,強迫 agent 改用我們的 test_model
    with currency_agent.override(model=test_model):
        result = currency_agent.run_sync("請問現在美金匯率多少?")
        
        # 斷言測試 (Assertions)
        # 1. 檢查工具是否有被呼叫
        assert test_model.tool_calls[0].tool_name == 'get_exchange_rate'
        assert test_model.tool_calls[0].args == {'currency_code': 'USD'}
        
        # 2. 檢查結果是否如預期 (32.5 是工具回傳的值,被當成最終解答)
        assert '32.5' in result.data

在這個測試中,我們完全沒有發送任何網路請求給 OpenAI,測試執行時間不到 0.01 秒,而且結果是 100% 確定性的 (Deterministic)。

覆寫依賴項 (Overriding Dependencies)

在先前的章節中,我們學過使用 deps 注入依賴項 (如資料庫連線)。在撰寫測試時,依賴注入架構的優勢就完全體現出來了:你只需要在測試時注入一個假的 (Mock) 資料庫實例即可。

# 在真實環境中執行
result = agent.run_sync("查詢餘額", deps=ProductionDatabase())

# 在測試環境中執行
def test_agent_with_mock_db():
    fake_db = MockDatabase(balance=1000)
    result = agent.run_sync("查詢餘額", deps=fake_db)
    # 斷言結果

如果你是在較大範圍的整合測試中,不方便直接修改 run_sync 的參數,Pydantic AI 同樣提供了 .override(deps=...) 的語法,讓你在上下文管理器內全域替換掉 Agent 的依賴項。

透過 TestModel 模擬 LLM 以及彈性的 Dependency Overrides,你可以輕易地為你的 Pydantic AI 專案建立高覆蓋率的自動化測試,確保應用程式的穩定運作。