summaryrefslogtreecommitdiff
path: root/factory
diff options
context:
space:
mode:
authorRaphaël Barrois <raphael.barrois@polytechnique.org>2013-03-04 23:18:02 +0100
committerRaphaël Barrois <raphael.barrois@polytechnique.org>2013-03-04 23:18:02 +0100
commit7d792430e103984a91c102c33da79be2426bc632 (patch)
tree0044c449e933e93d4fecf1e80e6b361c6c3e8150 /factory
parent9422cf12516143650f1014f34f996260c00d4c0a (diff)
downloadfactory-boy-7d792430e103984a91c102c33da79be2426bc632.tar
factory-boy-7d792430e103984a91c102c33da79be2426bc632.tar.gz
Add a 'after post_generation' hook to Factory.
Use it in DjangoModelFactory to save objects again if a post_generation hook ran. Signed-off-by: Raphaël Barrois <raphael.barrois@polytechnique.org>
Diffstat (limited to 'factory')
-rw-r--r--factory/base.py24
-rw-r--r--factory/declarations.py2
2 files changed, 24 insertions, 2 deletions
diff --git a/factory/base.py b/factory/base.py
index 28d7cdb..3ebc746 100644
--- a/factory/base.py
+++ b/factory/base.py
@@ -616,12 +616,27 @@ class Factory(BaseFactory):
obj = cls._prepare(create, **attrs)
# Handle post-generation attributes
+ results = {}
for name, decl in sorted(postgen_declarations.items()):
extracted, extracted_kwargs = postgen_attributes[name]
- decl.call(obj, create, extracted, **extracted_kwargs)
+ results[name] = decl.call(obj, create, extracted, **extracted_kwargs)
+
+ cls._after_postgeneration(obj, create, results)
+
return obj
@classmethod
+ def _after_postgeneration(cls, obj, create, results=None):
+ """Hook called after post-generation declarations have been handled.
+
+ Args:
+ obj (object): the generated object
+ create (bool): whether the strategy was 'build' or 'create'
+ results (dict or None): result of post-generation declarations
+ """
+ pass
+
+ @classmethod
def build(cls, **kwargs):
attrs = cls.attributes(create=False, extra=kwargs)
return cls._generate(False, attrs)
@@ -657,6 +672,13 @@ class DjangoModelFactory(Factory):
"""Create an instance of the model, and save it to the database."""
return target_class._default_manager.create(*args, **kwargs)
+ @classmethod
+ def _after_postgeneration(cls, obj, create, results=None):
+ """Save again the instance if creating and at least one hook ran."""
+ if create and results:
+ # Some post-generation hooks ran, and may have modified us.
+ obj.save()
+
class MogoFactory(Factory):
"""Factory for mogo objects."""
diff --git a/factory/declarations.py b/factory/declarations.py
index d3d7659..366c2c8 100644
--- a/factory/declarations.py
+++ b/factory/declarations.py
@@ -412,7 +412,7 @@ class PostGeneration(PostGenerationDeclaration):
self.function = function
def call(self, obj, create, extracted=None, **kwargs):
- self.function(obj, create, extracted, **kwargs)
+ return self.function(obj, create, extracted, **kwargs)
def post_generation(*args, **kwargs):