如何在单元测试中优雅修补Pydantic V2模型字段？

我正尝试在单元测试中修补Pydantic V2+模型的一个或多个字段。

需求背景 #

我想用更精简的枚举模拟某个枚举字段，减少测试断言的冗余信息。例如，我有如下代码：

from pydantic import BaseModel

class Foo(BaseModel):
    field: FooEnum

其中FooEnum定义如下：

from enum import Enum

class FooEnum(Enum):
    """包含大量字段的枚举"""

当我需要验证向Foo传入无效枚举的行为时，错误信息会非常繁琐。因此我希望用更小的MockFooEnum来模拟Foo.field，提升测试可读性，同时避免原枚举新增字段时需要同步修改单元测试。

以下方案可以运行，但它是通过研究Pydantic源码后拼凑出来的：

from contextlib import contextmanager
from unittest import mock
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from pydantic import BaseModel
    from pydantic.fields import FieldInfo


@contextmanager
def patch_pydantic_model_field(
    target: type[BaseModel], field_overrides: dict[str, FieldInfo]
) -> Generator[type[BaseModel], None, None]:
    model_fields = target.model_fields

    with mock.patch.object(
        target=target, attribute="model_fields", new_callable=mock.PropertyMock
    ) as mock_fields:
        # 用新的模拟字段覆盖模型
        mock_fields.return_value = model_fields | field_overrides
        target.model_rebuild(force=True)

        yield target

    target.model_rebuild(force=True)

使用示例 #

@pytest.mark.parametrize(
    "field_overrides",
    [{"field": FieldInfo(annotation=MockFooEnum, required=True)}],
)
def test_foo(field_overrides):
    with patch_pydantic_model_field(Foo, field_overrides):
        assert <something_with_mocked_model>

这个方法看起来比较粗糙，我想知道是否有更优雅的实现方式。

方案一：动态创建模型子类（推荐） #

直接修改原模型容易产生测试污染，更稳妥的方式是基于原模型动态创建子类，仅在子类中覆盖目标字段，完全隔离原模型的影响。

from typing import Type, TypeVar
from pydantic import BaseModel, FieldInfo

T = TypeVar('T', bound=BaseModel)

def create_mocked_model(original_model: Type[T], field_overrides: dict[str, FieldInfo]) -> Type[T]:
    new_fields = {}
    # 遍历原模型字段，保留未覆盖的，替换需要修改的
    for name, field in original_model.model_fields.items():
        if name in field_overrides:
            override = field_overrides[name]
            new_fields[name] = (override.annotation, override)
        else:
            new_fields[name] = (field.annotation, field)
    
    # 动态生成带Mock字段的子类
    return type(f"Mocked{original_model.__name__}", (original_model,), new_fields)

使用示例

@pytest.mark.parametrize("invalid_value", ["invalid", 123])
def test_foo_invalid_enum(invalid_value):
    # 创建仅使用MockFooEnum的模型子类
    MockedFoo = create_mocked_model(Foo, {"field": FieldInfo(annotation=MockFooEnum, required=True)})
    
    with pytest.raises(ValueError) as exc_info:
        MockedFoo(field=invalid_value)
    
    # 断言简洁的错误信息
    permitted_values = list(MockFooEnum.__members__.keys())
    assert f"value is not a valid enumeration member; permitted: {permitted_values}" in str(exc_info.value)

优势

• 无测试污染：不会修改原模型的任何属性，测试之间完全独立
• 兼容性强：基于Pydantic的继承特性实现，不依赖内部私有属性
• 可读性高：代码逻辑清晰，一眼就能看出是创建了测试专用模型

方案二：优化后的上下文管理器 #

如果确实需要临时修改原模型，可以简化原方案，去掉不必要的Mock操作，直接操作model_fields：

from contextlib import contextmanager
from typing import Type, TypeVar, Generator
from pydantic import BaseModel, FieldInfo

T = TypeVar('T', bound=BaseModel)

@contextmanager
def patch_pydantic_model_field(original_model: Type[T], field_overrides: dict[str, FieldInfo]) -> Generator[Type[T], None, None]:
    # 保存原字段快照
    original_fields = {name: original_model.model_fields[name] for name in field_overrides}
    
    try:
        # 更新字段并重建模型
        for name, override in field_overrides.items():
            original_model.model_fields[name] = override
        original_model.model_rebuild(force=True)
        
        yield original_model
    finally:
        # 恢复原字段并还原模型
        for name, field in original_fields.items():
            original_model.model_fields[name] = field
        original_model.model_rebuild(force=True)

这个版本去掉了对mock.patch的依赖，直接操作模型的model_fields属性，逻辑更简洁，也减少了对Pydantic内部实现的依赖。

内容的提问来源于stack exchange，提问作者Josmoor98