Merge branch 'release' into extend-tables
This commit is contained in:
Cayo Puigdefabregas
commit
74ae5c87f5
|
|
@ -0,0 +1,293 @@
|
|||
import random
|
||||
from typing import Type, Callable, List
|
||||
|
||||
from django.core.management.base import BaseCommand
|
||||
from django.db import models, IntegrityError
|
||||
|
||||
from idhub.models import Event, Rol, Service, UserRol
|
||||
from idhub_auth.models import User
|
||||
|
||||
from faker import Faker
|
||||
|
||||
DEFAULT_OBJECTS_CREATED: int = 30
|
||||
RANDOM_STRING_LENGTH: int = 30
|
||||
EMAIL_RANDOM_STRING_LENGTH: int = 10
|
||||
EXISTING_EVENTS: int = 30
|
||||
|
||||
fake = Faker()
|
||||
|
||||
|
||||
class Command(BaseCommand):
|
||||
"""
|
||||
Populate the database with dummy values.
|
||||
You can either specify which model to create objects for,
|
||||
or create objects for all models.
|
||||
If no data is specified it will create 30 events, users,
|
||||
services, roles, and user roles.
|
||||
"""
|
||||
|
||||
help = 'Populate the database with dummy values for testing.'
|
||||
|
||||
def __init__(self, *args, **kwargs):
|
||||
"""
|
||||
In the context of a Django management command,
|
||||
initializing lists like created_users in the constructor ensures
|
||||
that each run of the command has its own set of users, services,
|
||||
roles, etc., and does not unintentionally share or retain state
|
||||
across different invocations of the command.
|
||||
"""
|
||||
super().__init__(*args, **kwargs)
|
||||
self.created_users = []
|
||||
self.created_services = []
|
||||
self.created_roles = []
|
||||
|
||||
def handle(self, *args, **options):
|
||||
any_option_used = self.process_options(options)
|
||||
|
||||
if not any_option_used:
|
||||
self.create_all()
|
||||
|
||||
def process_options(self, input_options):
|
||||
option_methods = {
|
||||
'events': self.create_events,
|
||||
'users': self.create_users,
|
||||
'superusers': self.create_superusers,
|
||||
'services': self.create_services,
|
||||
'roles': self.create_roles,
|
||||
'userroles': self.create_user_roles,
|
||||
'userrole': self.create_specific_user_role,
|
||||
'register': self.register_user,
|
||||
'register_superuser': self.register_superuser,
|
||||
'amount': self.create_all
|
||||
}
|
||||
|
||||
any_option_used = self.match_input_to_options(input_options,
|
||||
option_methods)
|
||||
|
||||
return any_option_used
|
||||
|
||||
def match_input_to_options(self, input_options, option_methods):
|
||||
any_option_used = False
|
||||
|
||||
for option, method in option_methods.items():
|
||||
is_valid_option = option in input_options and input_options[
|
||||
option] is not None
|
||||
if is_valid_option:
|
||||
self.call_selected_method(input_options, method, option)
|
||||
any_option_used = True
|
||||
|
||||
return any_option_used
|
||||
|
||||
def call_selected_method(self, input_options, method, option):
|
||||
args = self.create_argument_list(input_options, option)
|
||||
method(*args)
|
||||
|
||||
def create_argument_list(self, input_options, option):
|
||||
if isinstance(input_options[option], list):
|
||||
args = input_options[option]
|
||||
else:
|
||||
args = [input_options[option]]
|
||||
return args
|
||||
|
||||
def add_arguments(self, parser):
|
||||
parser.add_argument(
|
||||
'--amount', type=int, action='store',
|
||||
help='Create N objects (includes events, users, services, roles, and user roles)'
|
||||
)
|
||||
parser.add_argument('--events', type=int,
|
||||
help='Create the specified number of events')
|
||||
parser.add_argument('--users', type=int,
|
||||
help='Create the specified number of users')
|
||||
parser.add_argument('--superusers', type=int,
|
||||
help='Create the specified number of superusers')
|
||||
parser.add_argument('--services', type=int,
|
||||
help='Create the specified number of services')
|
||||
parser.add_argument('--roles', type=int,
|
||||
help='Create the specified number of roles')
|
||||
parser.add_argument('--userroles', type=int,
|
||||
help='Create the specified number of user roles')
|
||||
parser.add_argument(
|
||||
'--userrole', nargs=2, type=str,
|
||||
help='Create a user role for user U and service S',
|
||||
metavar=('U', 'S'),
|
||||
)
|
||||
parser.add_argument('--register', nargs=2, type=str,
|
||||
help='Create a user with email EMAIL and password PW',
|
||||
metavar=('EMAIL', 'PW'))
|
||||
parser.add_argument('--register-superuser', nargs=2, type=str,
|
||||
help='Create a superuser with email EMAIL and '
|
||||
'password PW',
|
||||
metavar=('EMAIL', 'PW'))
|
||||
|
||||
def register_user(self, email: str, password: str):
|
||||
"""
|
||||
Register a new user with the given email and password from the command line.
|
||||
"""
|
||||
try:
|
||||
self.create_user(email, password)
|
||||
self.stdout.write(f"Successfully registered user: {email}")
|
||||
except IntegrityError:
|
||||
self.stdout.write(f"Failed to register user: {email}")
|
||||
|
||||
def register_superuser(self, email: str, password: str):
|
||||
"""
|
||||
Register a new superuser with the given email and password from the command line.
|
||||
"""
|
||||
try:
|
||||
self.create_superuser(email, password)
|
||||
self.stdout.write(f"Successfully registered superuser: {email}")
|
||||
except IntegrityError:
|
||||
self.stdout.write(f"Failed to register superuser: {email}")
|
||||
|
||||
def create_all(self, amount=DEFAULT_OBJECTS_CREATED):
|
||||
self.create_events(amount)
|
||||
self.create_users(amount)
|
||||
self.create_roles(amount)
|
||||
self.create_services(amount)
|
||||
self.create_user_roles(amount)
|
||||
|
||||
def create_objects(self, model: Type[models.Model],
|
||||
data_generator: Callable, amount: int) -> List[
|
||||
models.Model]:
|
||||
"""
|
||||
Generic method to create objects for a given model and keep track of them.
|
||||
|
||||
Args:
|
||||
model: The Django model class for which objects are to be created.
|
||||
data_generator: A function that returns a dictionary of attributes for creating a model instance.
|
||||
amount: The number of objects to create.
|
||||
|
||||
Returns:
|
||||
A list of successfully created model instances.
|
||||
"""
|
||||
created_objects = []
|
||||
for _ in range(amount):
|
||||
try:
|
||||
model_instance = self.create_from_data(data_generator, model)
|
||||
created_objects.append(model_instance)
|
||||
except IntegrityError:
|
||||
self.print_failure_message(model)
|
||||
|
||||
self.print_amount_created(created_objects, model)
|
||||
return created_objects
|
||||
|
||||
def create_users(self, amount: int):
|
||||
def user_data():
|
||||
return {
|
||||
'email': fake.email(),
|
||||
'first_name': fake.first_name(),
|
||||
'last_name': fake.last_name()
|
||||
}
|
||||
created_users = self.create_objects(User, user_data, amount)
|
||||
self.created_users.extend(created_users)
|
||||
|
||||
def create_events(self, amount: int, user=None):
|
||||
def event_data():
|
||||
return {
|
||||
'type': random.randint(1, EXISTING_EVENTS),
|
||||
'message': fake.paragraph(nb_sentences=3),
|
||||
'user': user
|
||||
}
|
||||
self.create_objects(Event, event_data, amount)
|
||||
|
||||
def create_superusers(self, amount=0):
|
||||
"""Superusers can only be created from the specific command"""
|
||||
created_superuser_amount = 0
|
||||
for value in range(0, amount):
|
||||
email = fake.email()
|
||||
try:
|
||||
User.objects.create_superuser(email)
|
||||
created_superuser_amount += 1
|
||||
except IntegrityError:
|
||||
self.stdout.write("Couldn't create superuser " + email)
|
||||
self.stdout.write(f"Created {created_superuser_amount} users")
|
||||
|
||||
def create_services(self, amount: int):
|
||||
def service_data():
|
||||
domain = fake.text(max_nb_chars=200)
|
||||
description = fake.text(max_nb_chars=250)
|
||||
return {
|
||||
'domain': domain,
|
||||
'description': description,
|
||||
}
|
||||
|
||||
created_services = self.create_objects(Service, service_data, amount)
|
||||
self.associate_random_roles(created_services)
|
||||
self.created_services.extend(created_services)
|
||||
|
||||
def create_roles(self, amount: int):
|
||||
def role_data():
|
||||
name = fake.job()
|
||||
description = fake.text(max_nb_chars=250)
|
||||
return {'name': name, 'description': description}
|
||||
|
||||
created_roles = self.create_objects(Rol, role_data, amount)
|
||||
self.created_roles.extend(created_roles)
|
||||
|
||||
def create_user_roles(self, amount: int, user_id=None, service_id=None):
|
||||
def user_role_data():
|
||||
user = self.get_or_create_user(user_id)
|
||||
service = self.get_or_create_service(service_id)
|
||||
return {"user": user, "service": service}
|
||||
|
||||
self.create_objects(UserRol, user_role_data, amount)
|
||||
|
||||
def create_specific_user_role(self, user, service):
|
||||
self.create_user_roles(1, user, service)
|
||||
|
||||
def create_user(self, email: str, password: str):
|
||||
User.objects.create_user(email, password)
|
||||
|
||||
def create_superuser(self, email: str, password: str):
|
||||
User.objects.create_superuser(email, password)
|
||||
|
||||
def print_failure_message(self, model):
|
||||
self.stdout.write(f"Couldn't create {model.__name__} object.")
|
||||
|
||||
def print_amount_created(self, created_objects, model):
|
||||
self.stdout.write(f"Created {len(created_objects)} "
|
||||
f"{model.__name__} objects")
|
||||
|
||||
def create_from_data(self, data_generator, model):
|
||||
model_instance = model.objects.create(**data_generator())
|
||||
return model_instance
|
||||
|
||||
def associate_random_roles(self, created_services):
|
||||
for service in created_services:
|
||||
self.associate_service_to_random_roles(service)
|
||||
|
||||
def associate_service_to_random_roles(self, service):
|
||||
associated_roles = [self.get_or_create_role() for _ in range(
|
||||
random.randint(0, 2))]
|
||||
service.rol.set(associated_roles)
|
||||
|
||||
def get_or_create_user(self, user_id):
|
||||
if user_id is not None:
|
||||
user = User.objects.get(user_id=user_id)
|
||||
elif len(self.created_users) != 0:
|
||||
user = random.choice(self.created_users)
|
||||
else:
|
||||
self.create_users(1)
|
||||
user = self.created_users[0]
|
||||
|
||||
return user
|
||||
|
||||
def get_or_create_service(self, service_id):
|
||||
if service_id is not None:
|
||||
service = Service.objects.get(service_id=service_id)
|
||||
elif len(self.created_services) != 0:
|
||||
service = random.choice(self.created_services)
|
||||
else:
|
||||
self.create_services(1)
|
||||
service = self.created_services[0]
|
||||
|
||||
return service
|
||||
|
||||
def get_or_create_role(self):
|
||||
if len(self.created_roles) != 0:
|
||||
role = random.choice(self.created_roles)
|
||||
else:
|
||||
self.create_roles(1)
|
||||
role = self.created_roles[0]
|
||||
|
||||
return role
|
||||
|
|
@ -13,3 +13,4 @@ jsonref==1.1.0
|
|||
pyld==2.0.3
|
||||
more-itertools==10.1.0
|
||||
dj-database-url==2.1.0
|
||||
faker==21.0.0
|
||||
|
|
|
|||
Loading…
Reference in New Issue