funkwhale/api/funkwhale_api/users/factories.py

import factory
from django.contrib.auth.models import Permission
from django.utils import timezone

from funkwhale_api.factories import ManyToManyFromList, registry

from . import models


@registry.register
class GroupFactory(factory.django.DjangoModelFactory):
    name = factory.Sequence(lambda n: "group-{0}".format(n))

    class Meta:
        model = "auth.Group"

    @factory.post_generation
    def perms(self, create, extracted, **kwargs):
        if not create:
            # Simple build, do nothing.
            return

        if extracted:
            perms = [
                Permission.objects.get(
                    content_type__app_label=p.split(".")[0], codename=p.split(".")[1]
                )
                for p in extracted
            ]
            # A list of permissions were passed in, use them
            self.permissions.add(*perms)


@registry.register
class InvitationFactory(factory.django.DjangoModelFactory):
    owner = factory.LazyFunction(lambda: UserFactory())

    class Meta:
        model = "users.Invitation"

    class Params:
        expired = factory.Trait(expiration_date=factory.LazyFunction(timezone.now))


@registry.register
class UserFactory(factory.django.DjangoModelFactory):
    username = factory.Sequence(lambda n: "user-{0}".format(n))
    email = factory.Sequence(lambda n: "user-{0}@example.com".format(n))
    password = factory.PostGenerationMethodCall("set_password", "test")
    subsonic_api_token = None
    groups = ManyToManyFromList("groups")
    avatar = factory.django.ImageField()

    class Meta:
        model = "users.User"
        django_get_or_create = ("username",)

    class Params:
        invited = factory.Trait(invitation=factory.SubFactory(InvitationFactory))

    @factory.post_generation
    def perms(self, create, extracted, **kwargs):
        if not create:
            # Simple build, do nothing.
            return

        if extracted:
            perms = [
                Permission.objects.get(
                    content_type__app_label=p.split(".")[0], codename=p.split(".")[1]
                )
                for p in extracted
            ]
            # A list of permissions were passed in, use them
            self.user_permissions.add(*perms)

    @factory.post_generation
    def with_actor(self, create, extracted, **kwargs):
        if not create or not extracted:
            return
        self.actor = models.create_actor(self)
        self.save(update_fields=["actor"])
        return self.actor


@registry.register(name="users.SuperUser")
class SuperUserFactory(UserFactory):
    is_staff = True
    is_superuser = True