Pinry/pinry/core/tests/helpers.py

from django.conf import settings
from django.contrib.auth.models import Permission
from django.core.files import File

import factory
from taggit.models import Tag

from ..models import Image, Pin
from ...users.models import User


TEST_IMAGE_PATH = settings.SITE_ROOT + 'screenshot.png'


class UserFactory(factory.Factory):
    FACTORY_FOR = User

    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


class TagFactory(factory.Factory):
    FACTORY_FOR = Tag

    name = factory.Sequence(lambda n: 'tag_{}'.format(n))


class ImageFactory(factory.Factory):
    FACTORY_FOR = Image

    image = factory.LazyAttribute(lambda a: File(open(TEST_IMAGE_PATH)))


class PinFactory(factory.Factory):
    FACTORY_FOR = Pin

    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