48 lines
1.2 KiB
Python
48 lines
1.2 KiB
Python
import factory
|
|
from wagtail.wagtailcore.models import Page
|
|
from personalisation.models import PersonalisablePage
|
|
|
|
from tests.sandbox.pages.models import HomePage
|
|
|
|
class PersonalisablePageFactory(factory.DjangoModelFactory):
|
|
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 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'])
|