Merge branch 'release' into extend-tables

This commit is contained in:
Cayo Puigdefabregas 2024-01-03 10:17:49 +01:00
commit 74ae5c87f5
2 changed files with 294 additions and 0 deletions

View File

@ -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

View File

@ -13,3 +13,4 @@ jsonref==1.1.0
pyld==2.0.3 pyld==2.0.3
more-itertools==10.1.0 more-itertools==10.1.0
dj-database-url==2.1.0 dj-database-url==2.1.0
faker==21.0.0