33 lines
917 B
Python
33 lines
917 B
Python
import random
|
|
|
|
import factory
|
|
|
|
from user.models import UserGroup
|
|
|
|
class_types = ['DA', 'KV', 'INF', 'EE']
|
|
class_suffix = ['A', 'B', 'C', 'D', 'E']
|
|
|
|
|
|
# TODO: refactor to use a trait synching the year with the infix year
|
|
# TODO: refactor to have non-overlapping user groups?
|
|
|
|
|
|
class UserGroupFactory(factory.django.DjangoModelFactory):
|
|
class Meta:
|
|
model = UserGroup
|
|
|
|
name = factory.Sequence(lambda n: '{}{}{}'.format(random.choice(class_types), '18', class_suffix[n % len(class_suffix)]))
|
|
year = factory.LazyAttribute(lambda x: random.choice([2017, 2018, 2019]))
|
|
is_deleted = False
|
|
|
|
@factory.post_generation
|
|
def users(self, create, extracted, **kwargs):
|
|
if not create:
|
|
# Simple build, do nothing.
|
|
return
|
|
|
|
if extracted:
|
|
# A list of groups were passed in, use them
|
|
for user in extracted:
|
|
self.users.add(user)
|