Improve handling of custom build/create functions.

Only public attributes (i.e not starting with _) will be handled properly.

Signed-off-by: Raphaël Barrois <raphael.barrois@polyconseil.fr>

1 files changed, 17 insertions, 0 deletions

diff --git a/tests.py b/tests.py
index 89a9363..bc534a6 100644
--- a/tests.py
+++ b/tests.py
@@ -258,6 +258,23 @@ class FactoryCreationTestCase(unittest.TestCase):
 
         self.assertEqual(TestFactory.default_strategy, STUB_STRATEGY)
 
+    def testCustomCreation(self):
+        class TestModelFactory(Factory):
+            @classmethod
+            def _prepare(cls, create, **kwargs):
+                kwargs['four'] = 4
+                return super(TestModelFactory, cls)._prepare(create, **kwargs)
+
+        b = TestModelFactory.build(one=1)
+        self.assertEqual(1, b.one)
+        self.assertEqual(4, b.four)
+        self.assertEqual(None, b.id)
+
+        c = TestModelFactory(one=1)
+        self.assertEqual(1, c.one)
+        self.assertEqual(4, c.four)
+        self.assertEqual(1, c.id)
+
     # Errors
 
     def testNoAssociatedClassWithAutodiscovery(self):