I'm attempting to patch one or more fields of a pydantic model (v2+) in a unit test.
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.
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)
@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.
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()
If you love us? You can donate to us via Paypal or buy me a coffee so we can maintain and grow! Thank you!
Donate Us With