Merge pull request 'Added custom command to create example data' (#92) from example-data into release

Reviewed-on: https://gitea.pangea.org/trustchain-oc1-orchestral/IdHub/pulls/92

idhub/management/commands/create_example_data.py

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

requirements.txt

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