Logo Questions Linux Laravel Mysql Ubuntu Git Menu
 

Pydantic V2 patching model fields

Problem

I'm attempting to patch one or more fields of a pydantic model (v2+) in a unit test.

Why?

I want to mock some enum field, with a simpler/smaller enum to reduce the test assertion noise. For example, if I have;

from pydantic import BaseModel

class Foo(BaseModel):
    field: FooEnum

where FooEnum is

from enum import Enum

class FooEnum(Enum):
    """Some enum with lot's of fields"""

If I want to explicitly validate the behaviour when some invalid enum is passed to Foo, the error message can be a nuisance to deal with. As such, I wanted to mock Foo.field with some smaller MockFooEnum, to improve testing readability/remove the need to update this unit test when a new field is added to FooEnum.

Attempted solutions

The following approach "works", but this was the result of hacking around in pydantic source code, to see how I could patch fields.

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:
        # ? Override the model with new mocked fields.
        mock_fields.return_value = model_fields | field_overrides
        target.model_rebuild(force=True)

        yield target

    target.model_rebuild(force=True)

Usage

@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>

This approach seems a bit of a bodge and was curious if there is a nicer way to achieve the above.

like image 658
Josmoor98 Avatar asked Mar 09 '26 22:03

Josmoor98


1 Answers

I do not think using a mock value is a good idea. For testing you usually want to test as much of the actual code as possible.

If you know exactly what you are doing, you could alternative create a new modified model using the create_model function:

from enum import Enum
from pydantic import BaseModel, create_model

class FooEnumLarge(Enum):
    """Some enum with lot's of fields"""
    one = 1
    two = 2
    three = 3
    four = 4
    five = 5


class FooEnumSmall(Enum):
    """Some enum with fewer fields"""
    one = 1
    two = 2


class Foo(BaseModel):
    field: FooEnumLarge


FooTest = create_model("FooNew", __base__=Foo, field=(FooEnumSmall, ...))

print(Foo(field=1))
print(FooTest(field=1))

Which prints:

field=<FooEnumLarge.one: 1>
field=<FooEnumSmall.one: 1>

So you could write a context manager, that does not modify the existing, but returns a new class:

with context_overide_fields(Foo, ...) as FooNew:
    FooNew()
like image 188
Axel Donath Avatar answered Mar 12 '26 11:03

Axel Donath