tests: post_build test coverage

Author: iloveitaly

tests/test_factory_fields.py

@@ -166,6 +166,25 @@ def caption(cls, is_long: bool) -> str:
         assert result.caption == "just this"
 
 
+def test_post_build_classmethod() -> None:
+    @dataclass
+    class Model:
+        i: int
+        j: int
+
+    class Factory(DataclassFactory[Model]):
+        __model__ = Model
+
+        @classmethod
+        def post_build(cls, model: Model) -> Model:
+            model.i = model.j + 10
+            return model
+
+    result = Factory.build()
+
+    assert result.i == result.j + 10
+
+
 @pytest.mark.parametrize(
     "factory_field",
     [

tests/test_pydantic_factory.py

@@ -1100,3 +1100,20 @@ class PaymentFactory(ModelFactory[Payment]):
 
     instance = PaymentFactory.build(currency="DKK")
     assert instance.currency == "DKK"
+
+
+def test_post_build_classmethod() -> None:
+    class Model(BaseModel):
+        i: int
+        j: int
+
+    class Factory(ModelFactory[Model]):
+        __model__ = Model
+
+        @classmethod
+        def post_build(cls, model: Model) -> Model:
+            model.i = model.j + 10
+            return model
+
+    result = Factory.build()
+    assert result.i == result.j + 10

tests/test_type_coverage_generation.py

@@ -202,6 +202,26 @@ def i(cls, j: int) -> int:
     assert results[0].i == results[0].j + 10
 
 
+def test_coverage_post_build() -> None:
+    @dataclass
+    class Model:
+        i: int
+        j: int
+
+    class Factory(DataclassFactory[Model]):
+        __model__ = Model
+
+        @classmethod
+        def post_build(cls, model: Model) -> Model:
+            model.i = model.j + 10
+            return model
+
+    results = list(Factory.coverage())
+    assert len(results) == 1
+
+    assert results[0].i == results[0].j + 10
+
+
 class CustomInt:
     def __init__(self, value: int) -> None:
         self.value = value