단위 테스트
단위 테스트 (Unit testing)
Pydantic AI 코드에 대한 단위 테스트 작성은 다른 어떤 파이썬 코드에 대한 단위 테스트와 똑같아요. 대부분 새로운 것이 아니기 때문에, 이런 테스트를 작성하고 실행하는 데는 잘 정립된 도구와 패턴이 이미 있어요.
출처: 문서
본문
자신이 더 잘 안다고 정말 확신하지 않는 한, 대략 이 전략을 따르는 게 좋아요:
- 테스트 하니스로
pytest사용 - 긴 어서션을 입력하다 싶으면 inline-snapshot 사용
- 마찬가지로 dirty-equals는 큰 데이터 구조를 비교할 때 유용함
- 실제 모델 대신
TestModel또는FunctionModel을 사용해 실제 LLM 호출의 비용·지연·변동성을 피함 - 이미지를 생성하는 코드에는
TestImageGenerationModel을 같은 방식으로 사용 — 이미지 생성: 테스팅 참고 - 애플리케이션 로직 안에서 에이전트의 모델·의존성·toolsets을 교체하려면
Agent.override사용 - 실수로 비테스트 모델에 요청이 나가지 않도록
ALLOW_MODEL_REQUESTS=False를 전역으로 설정
TestModel로 단위 테스트 (Unit testing with TestModel)
애플리케이션 코드 대부분을 실행해 보는 가장 단순하고 빠른 방법은 TestModel을 쓰는 거예요. 기본적으로 에이전트의 모든 도구를 호출한 다음, 에이전트의 반환 타입에 따라 일반 텍스트 또는 구조화 응답을 반환합니다.
TestModel은 마법이 아니에요
TestModel의 "똑똑한"(하지만 너무 똑똑하진 않은) 부분은 등록된 도구의 스키마를 기반으로 함수 도구와 출력 타입에 대한 유효한 구조화 데이터를 생성하려 시도한다는 점이에요.
TestModel에는 ML이나 AI가 없어요. 단지 도구의 JSON 스키마를 만족하는 데이터를 생성하려는 평범한 절차적 파이썬 코드일 뿐이에요.
결과 데이터는 예쁘거나 관련성 있게 보이진 않겠지만, 대부분의 경우 Pydantic 검증을 통과할 거예요. 더 정교한 것을 원한다면 FunctionModel을 사용하고 자신만의 데이터 생성 로직을 작성하세요.
네이티브 도구로 에이전트 테스팅
TestModel은 프로바이더가 실행하는 네이티브 도구를 에뮬레이트할 수 없어요. 프로덕션 에이전트가 capabilities로 네이티브 도구를 구성했다면, 테스트가 네이티브 도구가 모델에 전달되는지를 구체적으로 확인하는 게 아니라면 테스트에서 agent.override(model=TestModel(), native_tools=[])로 그것을 오버라이드하세요.
다음 애플리케이션 코드에 대한 단위 테스트를 작성해 봅시다:
import asyncio
from datetime import date
from pydantic_ai import Agent, RunContext
from fake_database import DatabaseConn # (1)
from weather_service import WeatherService # (2)
weather_agent = Agent(
'openai:gpt-5.2',
deps_type=WeatherService,
instructions='Providing a weather forecast at the locations the user provides.',
)
@weather_agent.tool
def weather_forecast(
ctx: RunContext[WeatherService], location: str, forecast_date: date
) -> str:
if forecast_date < date.today(): # (3)
return ctx.deps.get_historic_weather(location, forecast_date)
else:
return ctx.deps.get_forecast(location, forecast_date)
async def run_weather_forecast( # (4)
user_prompts: list[tuple[str, int]], conn: DatabaseConn
):
"""Run weather forecast for a list of user prompts and save."""
async with WeatherService() as weather_service:
async def run_forecast(prompt: str, user_id: int):
result = await weather_agent.run(prompt, deps=weather_service)
await conn.store_forecast(user_id, result.output)
# run all prompts in parallel
await asyncio.gather(
*(run_forecast(prompt, user_id) for (prompt, user_id) in user_prompts)
)
DatabaseConn은 데이터베이스 연결을 담는 클래스예요.
WeatherService는 날씨 예보와 기상 이력 데이터를 가져오는 메서드를 가져요.
날짜가 과거인지 미래인지에 따라 다른 엔드포인트를 호출해야 해요. 이 세부사항이 왜 중요한지는 아래에서 볼 수 있어요.
이 함수는 에이전트와 함께 우리가 테스트하려는 코드예요.
여기에는 (user_prompt, user_id) 튜플 리스트를 받아 각 프롬프트에 대한 날씨 예보를 구하고 결과를 데이터베이스에 저장하는 함수가 있어요.
우리는 특정 객체를 목킹하거나 테스트 객체를 주입하려고 코드를 수정하지 않고 이 코드를 테스트하고 싶어요.
TestModel로 테스트를 작성하는 방법이에요:
from datetime import timezone
import pytest
from dirty_equals import IsNow, IsStr
from pydantic_ai import models, capture_run_messages, RequestUsage
from pydantic_ai.models.test import TestModel
from pydantic_ai import (
ModelResponse,
TextPart,
ToolCallPart,
ToolReturnPart,
UserPromptPart,
ModelRequest,
)
from fake_database import DatabaseConn
from weather_app import run_weather_forecast, weather_agent
pytestmark = pytest.mark.anyio # (1)
models.ALLOW_MODEL_REQUESTS = False # (2)
async def test_forecast():
conn = DatabaseConn()
user_id = 1
with capture_run_messages() as messages:
with weather_agent.override(model=TestModel()): # (3)
prompt = 'What will the weather be like in London on 2024-11-28?'
await run_weather_forecast([(prompt, user_id)], conn) # (4)
forecast = await conn.get_forecast(user_id)
assert forecast == '{"weather_forecast":"Sunny with a chance of rain"}' # (5)
assert messages == [ # (6)
ModelRequest(
parts=[
UserPromptPart(
content='What will the weather be like in London on 2024-11-28?',
timestamp=IsNow(tz=timezone.utc), # (7)
),
],
instructions='Providing a weather forecast at the locations the user provides.',
timestamp=IsNow(tz=timezone.utc),
run_id=IsStr(),
conversation_id=IsStr(),
),
ModelResponse(
parts=[
ToolCallPart(
tool_name='weather_forecast',
args={
'location': 'a',
'forecast_date': '2024-01-01', # (8)
},
tool_call_id=IsStr(),
)
],
usage=RequestUsage(
input_tokens=60,
output_tokens=7,
),
model_name='test',
timestamp=IsNow(tz=timezone.utc),
provider_name='test',
run_id=IsStr(),
conversation_id=IsStr(),
),
ModelRequest(
parts=[
ToolReturnPart(
tool_name='weather_forecast',
content='Sunny with a chance of rain',
tool_call_id=IsStr(),
timestamp=IsNow(tz=timezone.utc),
),
],
instructions='Providing a weather forecast at the locations the user provides.',
timestamp=IsNow(tz=timezone.utc),
run_id=IsStr(),
conversation_id=IsStr(),
),
ModelResponse(
parts=[
TextPart(
content='{"weather_forecast":"Sunny with a chance of rain"}',
)
],
usage=RequestUsage(
input_tokens=66,
output_tokens=16,
),
model_name='test',
timestamp=IsNow(tz=timezone.utc),
provider_name='test',
run_id=IsStr(),
conversation_id=IsStr(),
),
]
비동기 테스트를 실행하려고 anyio를 사용해요.
이것은 테스트 중 실수로 LLM에 실제 요청을 보내지 않도록 하는 안전 장치예요. 자세한 내용은 ALLOW_MODEL_REQUESTS를 참고하세요.
Agent.override를 사용해 에이전트의 모델을 TestModel로 교체해요. override의 좋은 점은 에이전트 run* 메서드 호출 지점에 접근할 필요 없이 에이전트 안의 모델을 교체할 수 있다는 거예요.
이제 override 컨텍스트 매니저 안에서 우리가 테스트하려는 함수를 호출해요.
기본적으로 TestModel은 수행된 도구 호출과 반환된 것을 요약한 JSON 문자열을 반환해요. 응답을 도메인에 더 가깝게 커스터마이즈하고 싶다면 TestModel을 정의할 때 custom_output_text='Sunny'를 추가할 수 있어요.
지금까지 우리는 실제로 어떤 도구가 어떤 값으로 호출됐는지 모르므로, capture_run_messages를 사용해 가장 최근 실행의 메시지를 검사하고 에이전트와 모델 사이의 교환이 예상대로 일어났는지 어서션할 수 있어요.
IsNow 헬퍼는 시간이 지나며 바뀌는 타임스탬프가 포함된 데이터에서도 선언적 어서션을 쓸 수 있게 해 줘요.
TestModel은 프롬프트에서 값을 추출하는 똑똑한 짓을 하지 않으므로, 이 값들은 하드코딩돼 있어요.
FunctionModel로 단위 테스트 (Unit testing with FunctionModel)
위의 테스트는 훌륭한 시작이지만, 세심한 독자라면 TestModel이 과거 날짜로 weather_forecast를 호출하므로 WeatherService.get_forecast는 전혀 호출되지 않는다는 걸 알아챌 거예요.
weather_forecast를 완전히 실행하려면 FunctionModel을 사용해 도구가 호출되는 방식을 커스터마이즈해야 해요.
FunctionModel로 커스텀 입력으로 weather_forecast 도구를 테스트하는 예제예요:
import re
import pytest
from pydantic_ai import models
from pydantic_ai import (
ModelMessage,
ModelResponse,
TextPart,
ToolCallPart,
)
from pydantic_ai.models.function import AgentInfo, FunctionModel
from fake_database import DatabaseConn
from weather_app import run_weather_forecast, weather_agent
pytestmark = pytest.mark.anyio
models.ALLOW_MODEL_REQUESTS = False
def call_weather_forecast( # (1)
messages: list[ModelMessage], info: AgentInfo
) -> ModelResponse:
if len(messages) == 1:
# first call, call the weather forecast tool
user_prompt = messages[0].parts[-1]
m = re.search(r'\d{4}-\d{2}-\d{2}', user_prompt.content)
assert m is not None
args = {'location': 'London', 'forecast_date': m.group()} # (2)
return ModelResponse(parts=[ToolCallPart('weather_forecast', args)])
else:
# second call, return the forecast
msg = messages[-1].parts[0]
assert msg.part_kind == 'tool-return'
return ModelResponse(parts=[TextPart(f'The forecast is: {msg.content}')])
async def test_forecast_future():
conn = DatabaseConn()
user_id = 1
with weather_agent.override(model=FunctionModel(call_weather_forecast)): # (3)
prompt = 'What will the weather be like in London on 2032-01-01?'
await run_weather_forecast([(prompt, user_id)], conn)
forecast = await conn.get_forecast(user_id)
assert forecast == 'The forecast is: Rainy with a chance of sun'
LLM 대신 FunctionModel이 호출할 함수 call_weather_forecast를 정의해요. 이 함수는 실행을 구성하는 ModelMessage 리스트와, 에이전트·함수 도구·반환 도구에 대한 정보를 담은 AgentInfo에 접근할 수 있어요.
우리 함수는 프롬프트에서 날짜를 추출하려 약간 지능적으로 동작하지만 위치는 그냥 하드코딩해요.
FunctionModel을 사용해 에이전트의 모델을 우리 커스텀 함수로 교체해요.
교체 모델이 요청 사이에 상태를 유지해야 한다면, FunctionModel은 함수 대신 async def __call__을 가진 호출 가능한 인스턴스도 받아요. 예제는 FunctionModel API 문서를 참고하세요.
pytest 픽스처로 모델 오버라이드 (Overriding model via pytest fixtures)
모두 모델 오버라이드가 필요한 테스트를 많이 작성한다면, pytest fixtures를 사용해 TestModel 또는 FunctionModel로 모델을 재사용 가능하게 오버라이드할 수 있어요.
TestModel로 모델을 오버라이드하는 픽스처 예제예요:
import pytest
from pydantic_ai.models.test import TestModel
from weather_app import weather_agent
@pytest.fixture
def override_weather_agent():
with weather_agent.override(model=TestModel()):
yield
async def test_forecast(override_weather_agent: None):
...
# test code here
더 알아보기 (Learn more)
- pytest — 테스트 하니스.
TestModel— 빠른 테스트용 모델.FunctionModel— 커스텀 데이터 생성 로직.