Use wagtail-factories for testing

Makefile

@@ -10,7 +10,9 @@ clean:
 	find . -name '*.egg-info' -delete
 
 requirements:
-	pip install --upgrade -e .
+	pip install --upgrade -e .[test]
+
+install: develop
 
 develop: clean requirements
 

setup.py

@@ -13,6 +13,7 @@ tests_require = [
     'pytest-sugar==0.7.1',
     'freezegun==0.3.8',
     'factory_boy==2.8.1',
+    'wagtail_factories==0.2.0',
 ]
 
 setup(

tests/factories/page.py

@@ -2,49 +2,17 @@ from __future__ import absolute_import, unicode_literals
 
 import factory
 from wagtail.wagtailcore.models import Page
+from wagtail_factories import PageFactory
 
 from personalisation.models import PersonalisablePage
 from tests.sandbox.pages.models import HomePage
 
 
-class PersonalisablePageFactory(factory.DjangoModelFactory):
+class PersonalisablePageFactory(PageFactory):
     class Meta:
         model = PersonalisablePage
 
-    @classmethod
-    def _create(cls, *args, **kwargs):
-        obj = super(PersonalisablePageFactory, cls)._build(*args, **kwargs)
-        if not obj.title:
-            obj.title = "Page-Test"
-        return obj
 
-
-class HomePageFactory(factory.DjangoModelFactory):
+class HomePageFactory(PageFactory):
     class Meta:
         model = HomePage
-
-    @classmethod
-    def build(cls, *args, **kwargs):
-        obj = super(HomePageFactory, cls)._build(*args, **kwargs)
-
-        for part in ('subtitle', 'body'):
-            if not getattr(obj, part):
-                setattr(obj, part, "{}".format(part))
-        return obj
-
-
-class SiteRootFactory(factory.DjangoModelFactory):
-    title = 'site-root'
-    depth = 2
-
-    class Meta:
-        model = Page
-
-    @classmethod
-    def _create(cls, *args, **kwargs):
-        try:
-            root = Page.objects.get(depth=0)
-        except Page.DoesNotExist:
-            root = Page.add_root(title='root')
-
-        return root.add_child(title=kwargs['title'])

tests/factories/site.py

@@ -1,17 +0,0 @@
-from __future__ import absolute_import, unicode_literals
-
-import factory
-from wagtail.wagtailcore.models import Site
-
-from tests.factories.page import SiteRootFactory
-
-
-class SiteFactory(factory.DjangoModelFactory):
-    hostname = 'localhost'
-    port = 8000
-    site_name = 'Testing Site'
-    root_page = factory.SubFactory(SiteRootFactory)
-    is_default_site = True
-
-    class Meta:
-        model = Site

tests/unit/test_middleware.py

@@ -6,11 +6,11 @@ import pytest
 from django.test.client import Client
 from freezegun import freeze_time
 from wagtail.wagtailcore.models import Page
+from wagtail_factories import SiteFactory
 
 from tests.factories.segment import (
     QueryRuleFactory, ReferralRuleFactory, SegmentFactory, TimeRuleFactory,
     DayRuleFactory, VisitCountRuleFactory, DeviceRuleFactory)
-from tests.factories.site import SiteFactory
 
 
 @pytest.mark.django_db
@@ -20,7 +20,7 @@ class TestUserSegmenting(object):
         """
         Sets up a site root to test segmenting
         """
-        self.site = SiteFactory()
+        self.site = SiteFactory(is_default_site=True)
 
     def test_no_segments(self, client):
         request = client.get('/')
@@ -164,7 +164,7 @@ class TestUserSegmenting(object):
         segment = SegmentFactory(name='Visit Count')
 
         visit_count_rule = VisitCountRuleFactory(
-            counted_page=Page.objects.filter(slug="root").first(),
+            counted_page=self.site.root_page,
             segment=segment
         )
 
@@ -262,14 +262,14 @@ class TestUserSegmenting(object):
         )
 
         client.get('/?test=test')
-        
+
         assert any(item['encoded_name'] == 'match-any' for item in client.session['segments'])
 
 @pytest.mark.django_db
 class TestUserVisitCount(object):
 
     def setup(self):
-        self.site = SiteFactory()
+        self.site = SiteFactory(is_default_site=True)
 
         # TODO: Set up a bunch of pages for testing the visit count
 

tests/unit/test_models.py

@@ -6,9 +6,9 @@ import pytest
 from django.http import HttpRequest
 from freezegun import freeze_time
 from wagtail.wagtailcore.models import Site
+from wagtail_factories import SiteFactory
 
 from personalisation import models
-from tests.factories.site import SiteFactory
 
 """Time Rule tests"""
 @freeze_time("10:00:00")