From 7a85f1b5145f08bc86635bf1e664000e61940bc3 Mon Sep 17 00:00:00 2001
From: Krzysztof Klimonda <kklimonda@syntaxhighlighted.com>
Date: Mon, 4 Mar 2013 14:39:26 -0800
Subject: [PATCH] Use post_generation hooks to populate models with some data

Factory Boy provides us with a better way of populating complicated
fields (like M2M relations) than overriding _prepare class method,
so we should be using them.
---
 pinry/core/tests/__init__.py |  1 +
 pinry/core/tests/helpers.py  | 58 ++++++++++++++++++++++++++----------
 2 files changed, 44 insertions(+), 15 deletions(-)

diff --git a/pinry/core/tests/__init__.py b/pinry/core/tests/__init__.py
index 42ac7b9..a547926 100644
--- a/pinry/core/tests/__init__.py
+++ b/pinry/core/tests/__init__.py
@@ -1,3 +1,4 @@
 from .api import *
 from .forms import *
+from .helpers import PinFactoryTest
 from .views import *
\ No newline at end of file
diff --git a/pinry/core/tests/helpers.py b/pinry/core/tests/helpers.py
index 2ffe2cc..485a798 100644
--- a/pinry/core/tests/helpers.py
+++ b/pinry/core/tests/helpers.py
@@ -1,6 +1,8 @@
 from django.conf import settings
 from django.contrib.auth.models import Permission
 from django.core.files import File
+from django.db.models.query import QuerySet
+from django.test import TestCase
 
 import factory
 from taggit.models import Tag
@@ -18,17 +20,14 @@ class UserFactory(factory.Factory):
     username = factory.Sequence(lambda n: 'user_{}'.format(n))
     email = factory.Sequence(lambda n: 'user_{}@example.com'.format(n))
 
-    @classmethod
-    def _prepare(cls, create, **kwargs):
-        password = kwargs.pop('password', None)
-        user = super(UserFactory, cls)._prepare(create, **kwargs)
-        user.user_permissions = Permission.objects.filter(codename__in=['add_pin', 'add_image'])
-        if password:
-            user.set_password(password)
-            if create:
-                user.save()
-        return user
+    @factory.post_generation(extract_prefix='password')
+    def set_password(self, create, extracted, **kwargs):
+        self.set_password(extracted)
+        self.save()
 
+    @factory.post_generation(extract_prefix='user_permissions')
+    def set_user_permissions(self, create, extracted, **kwargs):
+        self.user_permissions = Permission.objects.filter(codename__in=['add_pin', 'add_image'])
 
 class TagFactory(factory.Factory):
     FACTORY_FOR = Tag
@@ -48,8 +47,37 @@ class PinFactory(factory.Factory):
     submitter = factory.SubFactory(UserFactory)
     image = factory.SubFactory(ImageFactory)
 
-    @classmethod
-    def _prepare(cls, create, **kwargs):
-        pin = super(PinFactory, cls)._prepare(create, **kwargs)
-        pin.tags.add(TagFactory())
-        return pin
\ No newline at end of file
+    @factory.post_generation(extract_prefix='tags')
+    def add_tags(self, create, extracted, **kwargs):
+        if isinstance(extracted, Tag):
+            self.tags.add(extracted)
+        elif isinstance(extracted, list):
+            self.tags.add(*extracted)
+        elif isinstance(extracted, QuerySet):
+            self.tags = extracted
+        else:
+            self.tags.add(TagFactory())
+
+
+class PinFactoryTest(TestCase):
+    def test_default_tags(self):
+        self.assertTrue(PinFactory().tags.get(pk=1).name.startswith('tag_'))
+
+    def test_custom_tag(self):
+        custom = 'custom_tag'
+        self.assertEqual(PinFactory(tags=Tag.objects.create(name=custom)).tags.get(pk=1).name, custom)
+
+    def test_custom_tags_list(self):
+        tags = TagFactory.create_batch(2)
+        PinFactory(tags=tags)
+        self.assertEqual(Tag.objects.count(), 2)
+
+    def test_custom_tags_queryset(self):
+        TagFactory.create_batch(2)
+        tags = Tag.objects.all()
+        PinFactory(tags=tags)
+        self.assertEqual(Tag.objects.count(), 2)
+
+    def test_empty_tags(self):
+        PinFactory(tags=[])
+        self.assertEqual(Tag.objects.count(), 0)
\ No newline at end of file