diff --git a/program/tests/factories.py b/program/tests/factories.py index 1ce0368c4d7268eb3943394f752c1cac8ff4a726..738c6ddcd28637f228ba27c85e41b3367405b2bc 100644 --- a/program/tests/factories.py +++ b/program/tests/factories.py @@ -1,15 +1,17 @@ from datetime import timedelta -from factory import Sequence +from factory import Sequence, post_generation from factory.django import DjangoModelFactory -from django.contrib.auth.models import User +from django.contrib.auth.models import Permission, User +from django.db.models import QuerySet from django.utils.timezone import now from program.models import ( FundingCategory, Host, Image, License, + Note, RRule, Schedule, Show, @@ -26,6 +28,26 @@ class CommonUserFactory(DjangoModelFactory): username = Sequence(lambda n: "common_%d" % n) +class UserWithPermissionsFactory(DjangoModelFactory): + class Meta: + model = User + + username = Sequence(lambda n: "user_with_perm%d" % n) + + @post_generation + def user_permissions(self, create, extracted, **kwargs): + if not create: + return + + if extracted: + if isinstance(extracted, Permission): + self.user_permissions.add(extracted) + + if isinstance(extracted, QuerySet): + for permission in extracted: + self.user_permissions.add(permission) + + class HostFactory(DjangoModelFactory): class Meta: model = Host @@ -90,3 +112,8 @@ class LicenseFactory(DjangoModelFactory): identifier = "pd" name = "Public Domain" + + +class NoteFactory(DjangoModelFactory): + class Meta: + model = Note