diff --git a/tests/factories/site.py b/tests/factories/site.py
new file mode 100644
index 0000000..9b012e1
--- /dev/null
+++ b/tests/factories/site.py
@@ -0,0 +1,25 @@
+import factory
+from django.utils.text import slugify
+from wagtail.wagtailcore.models import Site
+from wagtail_factories.factories import MP_NodeFactory
+
+from tests.sandbox.pages.models import HomePage
+
+
+class PageFactory(MP_NodeFactory):
+    title = 'Test page'
+    slug = factory.LazyAttribute(lambda obj: slugify(obj.title))
+
+    class Meta:
+        model = HomePage
+
+
+class SiteFactory(factory.DjangoModelFactory):
+    hostname = 'localhost'
+    port = factory.Sequence(lambda n: 81 + n)
+    site_name = 'Test site'
+    root_page = factory.SubFactory(PageFactory, parent=None)
+    is_default_site = False
+
+    class Meta:
+        model = Site
diff --git a/tests/sandbox/pages/migrations/0001_initial.py b/tests/sandbox/pages/migrations/0001_initial.py
index 56439e3..86ff359 100644
--- a/tests/sandbox/pages/migrations/0001_initial.py
+++ b/tests/sandbox/pages/migrations/0001_initial.py
@@ -1,10 +1,10 @@
 # -*- coding: utf-8 -*-
-# Generated by Django 1.11.1 on 2017-05-31 06:50
+# Generated by Django 1.11.1 on 2017-05-31 06:40
 from __future__ import unicode_literals
 
-from django.db import migrations, models
 import django.db.models.deletion
 import wagtail.wagtailcore.fields
+from django.db import migrations, models
 
 
 class Migration(migrations.Migration):
@@ -12,20 +12,23 @@ class Migration(migrations.Migration):
     initial = True
 
     dependencies = [
-        ('wagtail_personalisation', '0008_devicerule'),
+        ('wagtail_personalisation', '0009_auto_20170531_0428'),
     ]
 
     operations = [
         migrations.CreateModel(
             name='HomePage',
             fields=[
-                ('personalisablepage_ptr', models.OneToOneField(auto_created=True, on_delete=django.db.models.deletion.CASCADE, parent_link=True, primary_key=True, serialize=False, to='wagtail_personalisation.PersonalisablePage')),
-                ('subtitle', models.CharField(max_length=255)),
+                ('page_ptr', models.OneToOneField(auto_created=True, on_delete=django.db.models.deletion.CASCADE, parent_link=True, primary_key=True, serialize=False, to='wagtailcore.Page')),
+                ('is_segmented', models.BooleanField(default=False)),
+                ('subtitle', models.CharField(blank=True, default='', max_length=255)),
                 ('body', wagtail.wagtailcore.fields.RichTextField(blank=True, default='')),
+                ('canonical_page', models.ForeignKey(blank=True, null=True, on_delete=django.db.models.deletion.SET_NULL, related_name='variations', to='pages.HomePage')),
+                ('segment', models.ForeignKey(blank=True, null=True, on_delete=django.db.models.deletion.PROTECT, related_name='segments', to='wagtail_personalisation.Segment')),
             ],
             options={
                 'abstract': False,
             },
-            bases=('wagtail_personalisation.personalisablepage',),
+            bases=('wagtailcore.page', models.Model),
         ),
     ]
diff --git a/tests/sandbox/settings.py b/tests/sandbox/settings.py
index 19f2db3..bb2575d 100644
--- a/tests/sandbox/settings.py
+++ b/tests/sandbox/settings.py
@@ -63,7 +63,7 @@ MIDDLEWARE = (
 
 INSTALLED_APPS = (
     'wagtail_personalisation',
-    'tests',
+    'tests.sandbox.pages',
 
     'wagtail.contrib.modeladmin',
     'wagtail.wagtailsearch',
diff --git a/tests/unit/test_models.py b/tests/unit/test_models.py
index 07cbf66..50fdd55 100644
--- a/tests/unit/test_models.py
+++ b/tests/unit/test_models.py
@@ -6,7 +6,7 @@ import pytest
 from django.http import HttpRequest
 from freezegun import freeze_time
 from wagtail.wagtailcore.models import Site
-from wagtail_factories import SiteFactory
+from tests.factories.site import SiteFactory
 
 from wagtail_personalisation import rules