Create Test Cases to Insert Over 1,000 Records in Each DB Table to Optimize API Performance

server/cshr/management/commands/seeds.py

@@ -0,0 +1,108 @@
+import datetime
+from enum import Enum
+import os
+import django
+import random
+from django_seed import Seed
+from cshr.models.users import User, Office, UserSkills
+from cshr.models.vacations import Vacation
+from django.core.management.base import BaseCommand
+
+from cshr.management.commands.thanos import Thanos, ThanosWorlds
+
+
+os.environ.setdefault("DJANGO_SETTINGS_MODULE", "cshr.settings")
+django.setup()
+
+seeder = Seed.seeder()
+
+class Command(BaseCommand):
+    def add_arguments(self, parser):
+        parser.add_argument(
+            "--model",
+            nargs="+",
+            type=ThanosWorlds,
+            help="Should be `locations`, `vacations`, or `users`",
+        )
+        parser.add_argument(
+            "--number", nargs="+", type=int, help="The number of seeds to be created, by default its 500."
+        )
+
+    def handle(self, *args, **options):
+        """Use this command only in development mode."""
+        model = options["model"]
+        seeds_number = options["number"] or 500
+        if type(seeds_number) == list:
+            seeds_number = seeds_number[0]
+
+        if type(model) == list:
+            model = model[0]
+
+        if not model:
+            answer = input(
+                "Are you sure you want to run seeds over the whole app? Y/N "
+            )
+
+            if answer.lower() == "n":
+                return self.stdout.write(
+                    self.style.ERROR(
+                        "Aborted."
+                    )
+                )
+
+            if answer.lower() != "y":
+                return self.stdout.write(
+                    self.style.ERROR(
+                        "Wrong answer."
+                    )
+                )
+
+            model = "all"
+
+        Thanos.use(model, seeds_number)
+        return Thanos.fingers.snap()
+
+
+        # print("seeds_number", seeds_number)
+        # print("model", model)
+        # try:
+        #     seeder.add_entity(
+
+
+# 		Vacation,
+# 		1,
+# 		{
+# 			"applying_user": User.objects.get(id=random_id(1)),
+# 			"approval_user": User.objects.get(id=random_id(1)),
+# 		},
+# 	)
+#     inserted_pks = seeder.execute()
+#     self.stdout.write(
+# 		self.style.SUCCESS("Successfully created 1000 vacations.")
+# 	)
+# except Exception:
+#     self.stdout.write(
+#         self.style.ERROR(
+#         	"Failed to create vacation balances, check the records in the database if exist, or check if you migrated the database."
+#         )
+#     )
+
+
+# Seed Users without setting the many-to-many fields
+# seeder.add_entity(User, 50, {
+#     "location": lambda x: Office.objects.order_by('?').first(),  # Random Office
+# })
+
+
+# # Manually create many-to-many relationships after seeding
+# for user in User.objects.all():
+#     # Handle reporting_to relationship
+#     possible_reportings = User.objects.exclude(id=user.id)
+#     if possible_reportings.exists():
+#         user.reporting_to.set(random.sample(list(possible_reportings), k=random.randint(1, 5)))
+
+#     # Handle skills relationship
+#     possible_skills = UserSkills.objects.all()
+#     if possible_skills.exists():
+#         user.skills.set(random.sample(list(possible_skills), k=random.randint(1, 5)))
+#

server/cshr/management/commands/thanos.py

@@ -0,0 +1,84 @@
+from enum import Enum
+import os
+import django
+import random
+from django_seed import Seed
+from cshr.models.users import User, Office, UserSkills
+from cshr.models.vacations import Vacation
+
+os.environ.setdefault("DJANGO_SETTINGS_MODULE", "cshr.settings")
+django.setup()
+
+seeder = Seed.seeder()
+
+class ThanosWorlds(Enum):
+    LOCATIONS = "locations"
+    VACATIONS = "vacations"
+    USERS = "users"
+    ALL = "all"
+
+class ThanosFingers:
+    def __init__(self):
+        self.model = None
+        self.number = 0
+
+    def snap(self):
+        if self.model is None:
+            raise ValueError("Model not set. Use Thanos.worlds.use() to set the model.")
+
+        thanos = Thanos(self.model, self.number)
+        if self.model == ThanosWorlds.ALL.value:
+            return thanos.create_all()
+        else:
+            return thanos.create_individual(self.model)
+
+class Thanos:
+    worlds = None
+    fingers = ThanosFingers()
+
+    def __init__(self, model: ThanosWorlds, number: int):
+        self.model = model
+        self.number = number
+
+    @classmethod
+    def use(cls, model: ThanosWorlds, number: int = 1000):
+        cls.fingers.model = model
+        cls.fingers.number = number
+
+    def create_all(self):
+        seeder.add_entity(Office, self.number)
+        seeder.add_entity(User, self.number)
+        seeder.add_entity(Vacation, self.number)
+        seeder.execute()
+        return self.cry("All the necessary seeds have been successfully created.")
+
+    def create_individual(self, model: ThanosWorlds):
+        if model == ThanosWorlds.LOCATIONS:
+            seeder.add_entity(Office, self.number)
+        elif model == ThanosWorlds.VACATIONS:
+            seeder.add_entity(Vacation, self.number)
+        elif model == ThanosWorlds.USERS:
+            self.create_users()
+
+        inserted_pks = seeder.execute()
+        return self.cry(f"Seeded {self.number} {model.value}.")
+
+    def create_users(self):
+        locations_count = Office.objects.all().count()
+        init_len = 50
+        if locations_count < init_len:
+            self.create_locations(init=init_len)
+            return seeder.add_entity(User, init_len, {
+                "location": lambda x: Office.objects.order_by('?').first(),
+            })
+        return seeder.add_entity(User, self.number, {
+            "location": lambda x: Office.objects.order_by('?').first(),
+        })
+
+    def create_locations(self, init: int = None):
+        if init:
+            return seeder.add_entity(Office, init)
+        return seeder.add_entity(Office, self.number)
+
+    def cry(self, message: str):
+        print(message)

server/settings.py

@@ -31,7 +31,8 @@
     "django.contrib.staticfiles",
     "django_celery_beat",
     "cshr",
-    # Third party
+    # 3rd party
+    "django_seed",
     "corsheaders",
     "drf_yasg",
     "rest_framework",