From 8b1c2e9c020559024c1aa7192b0dfcdae83426cb Mon Sep 17 00:00:00 2001
From: Vanderson Mota <vanderson.mota@gmail.com>
Date: Sat, 13 Jun 2015 15:14:33 -0300
Subject: [PATCH] automatically allow unsaved instances for foreignkeys and
 OneToFields from django 1.8 up

---
 model_mommy/mommy.py               | 19 ++++++++++++-------
 requirements.txt                   |  1 +
 test/generic/tests/test_mommy.py   | 10 ++--------
 test/generic/tests/test_recipes.py |  5 +----
 4 files changed, 16 insertions(+), 19 deletions(-)

diff --git a/model_mommy/mommy.py b/model_mommy/mommy.py
index b998f0a..e6769fb 100644
--- a/model_mommy/mommy.py
+++ b/model_mommy/mommy.py
@@ -57,6 +57,7 @@ def validate_ipv6_address(v):
 from .exceptions import ModelNotFound, AmbiguousModelName, InvalidQuantityException, RecipeIteratorEmpty
 
 from six import string_types, advance_iterator, PY3
+from mock import patch
 
 recipes = None
 
@@ -91,7 +92,16 @@ def make(model, _quantity=None, make_m2m=False, **attrs):
     else:
         return mommy.make(**attrs)
 
+def _allow_unsaved(fn, *args, **kwargs):
+    def p(*args, **kwargs):
+        if django.VERSION >= (1, 8):
+            with patch.object(ForeignKey, 'allow_unsaved_instance_assignment', True):
+                with patch.object(OneToOneField, 'allow_unsaved_instance_assignment', True):
+                    return fn(*args, **kwargs)
+        return fn(*args, **kwargs)
+    return p
 
+@_allow_unsaved
 def prepare(model, _quantity=None, **attrs):
     """
     Creates BUT DOESN'T persist an instance from a given model its
@@ -108,7 +118,6 @@ def prepare(model, _quantity=None, **attrs):
     else:
         return mommy.prepare(**attrs)
 
-
 def _recipe(name):
     app, recipe_name = name.rsplit('.', 1)
     recipes = importlib.import_module('.'.join([app, 'mommy_recipes']))
@@ -285,12 +294,8 @@ def make(self, **attrs):
     def prepare(self, **attrs):
         '''Creates, but does not persist, an instance of the model
         associated with Mommy instance.'''
-        if django.VERSION >= (1, 8):
-            self.type_mapping[ForeignKey] = make
-            self.type_mapping[OneToOneField] = make
-        else:
-            self.type_mapping[ForeignKey] = prepare
-            self.type_mapping[OneToOneField] = prepare
+        self.type_mapping[ForeignKey] = prepare
+        self.type_mapping[OneToOneField] = prepare
         return self._make(commit=False, **attrs)
 
     def get_fields(self):
diff --git a/requirements.txt b/requirements.txt
index 54376ad..cdc24ed 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,2 +1,3 @@
 Django<1.9
 six
+mock
diff --git a/test/generic/tests/test_mommy.py b/test/generic/tests/test_mommy.py
index 8ae8525..ec835ac 100644
--- a/test/generic/tests/test_mommy.py
+++ b/test/generic/tests/test_mommy.py
@@ -196,10 +196,7 @@ def test_prepare_fk(self):
         self.assertIsInstance(dog, Dog)
         self.assertIsInstance(dog.owner, Person)
 
-        if django.VERSION >= (1, 8):
-            self.assertEqual(Person.objects.all().count(), 1)
-        else:
-            self.assertEqual(Person.objects.all().count(), 0)
+        self.assertEqual(Person.objects.all().count(), 0)
         self.assertEqual(Dog.objects.all().count(), 0)
 
     def test_prepare_one_to_one(self):
@@ -207,10 +204,7 @@ def test_prepare_one_to_one(self):
 
         self.assertEqual(LonelyPerson.objects.all().count(), 0)
         self.assertTrue(isinstance(lonely_person.only_friend, Person))
-        if django.VERSION >= (1, 8):
-            self.assertEqual(Person.objects.all().count(), 1)
-        else:
-            self.assertEqual(Person.objects.all().count(), 0)
+        self.assertEqual(Person.objects.all().count(), 0)
 
     def test_create_one_to_one(self):
         lonely_person = mommy.make(LonelyPerson)
diff --git a/test/generic/tests/test_recipes.py b/test/generic/tests/test_recipes.py
index 722628c..1dcd1e7 100644
--- a/test/generic/tests/test_recipes.py
+++ b/test/generic/tests/test_recipes.py
@@ -364,10 +364,7 @@ def test_do_query_lookup_empty_recipes(self):
         self.assertEqual(dog.owner.name, 'James')
 
         dog = dog_recipe.prepare(owner__name='Zezin')
-        if django.VERSION >= (1, 8):
-            self.assertEqual(Person.objects.count(), 2)
-        else:
-            self.assertEqual(Person.objects.count(), 1)
+        self.assertEqual(Person.objects.count(), 1)
         self.assertEqual(dog.owner.name, 'Zezin')
 
     def test_related_models_recipes(self):