summaryrefslogtreecommitdiff
path: root/factory/base.py
diff options
context:
space:
mode:
authorRaphaël Barrois <raphael.barrois@polytechnique.org>2013-04-02 23:34:28 +0200
committerRaphaël Barrois <raphael.barrois@polytechnique.org>2013-04-02 23:34:28 +0200
commit6532f25058a13e81b1365bb353848510821f571f (patch)
tree55a60a1568a84f06103197d7cf55b2e17f95fd71 /factory/base.py
parent69e7a86875f97dc12e941302fabe417122f2cb7e (diff)
downloadfactory-boy-6532f25058a13e81b1365bb353848510821f571f.tar
factory-boy-6532f25058a13e81b1365bb353848510821f571f.tar.gz
Add support for get_or_create in DjangoModelFactory.
Diffstat (limited to 'factory/base.py')
-rw-r--r--factory/base.py15
1 files changed, 14 insertions, 1 deletions
diff --git a/factory/base.py b/factory/base.py
index 58cd50b..71c2eb1 100644
--- a/factory/base.py
+++ b/factory/base.py
@@ -554,6 +554,7 @@ class DjangoModelFactory(Factory):
"""
ABSTRACT_FACTORY = True
+ FACTORY_DJANGO_GET_OR_CREATE = ()
@classmethod
def _get_manager(cls, target_class):
@@ -578,7 +579,19 @@ class DjangoModelFactory(Factory):
def _create(cls, target_class, *args, **kwargs):
"""Create an instance of the model, and save it to the database."""
manager = cls._get_manager(target_class)
- return manager.create(*args, **kwargs)
+
+ assert 'defaults' not in cls.FACTORY_DJANGO_GET_OR_CREATE, (
+ "'defaults' is a reserved keyword for get_or_create "
+ "(in %s.FACTORY_DJANGO_GET_OR_CREATE=%r)"
+ % (cls, cls.FACTORY_DJANGO_GET_OR_CREATE))
+
+ key_fields = {}
+ for field in cls.FACTORY_DJANGO_GET_OR_CREATE:
+ key_fields[field] = kwargs.pop(field)
+ key_fields['defaults'] = kwargs
+
+ obj, created = manager.get_or_create(*args, **key_fields)
+ return obj
@classmethod
def _after_postgeneration(cls, obj, create, results=None):