improved populate script now

1 files changed, 26 insertions, 5 deletions

diff --git a/ctrack/organisations/management/commands/populate_db.py b/ctrack/organisations/management/commands/populate_db.py
index 2dda06e..06759a9 100644
--- a/ctrack/organisations/management/commands/populate_db.py
+++ b/ctrack/organisations/management/commands/populate_db.py
@@ -1,8 +1,11 @@
+from random import randint
+
 from django.core.management import BaseCommand
 from django.core.management import CommandParser
 
 from ctrack.organisations.models import Mode
 from ctrack.organisations.models import Submode
+from ctrack.organisations.tests.factories import OrganisationFactory
 from ctrack.organisations.tests.factories import PersonFactory
 from ctrack.organisations.tests.factories import RoleFactory
 from ctrack.organisations.tests.factories import UserFactory
@@ -19,10 +22,12 @@ class Command(BaseCommand):
         parser.add_argument("number", nargs=1, type=int)
 
     def handle(self, *args, **options):
+        number = options["number"][0]
 
         # Set up some reasonable Modes and SubModes
         m1 = Mode.objects.create(descriptor="Rail")
         m2 = Mode.objects.create(descriptor="Maritime")
+
         sb1 = Submode.objects.create(descriptor="Light Rail", mode=m1)
         sb2 = Submode.objects.create(descriptor="Rail Maintenance", mode=m1)
         sb3 = Submode.objects.create(descriptor="Rail Infrastructure", mode=m1)
@@ -31,15 +36,31 @@ class Command(BaseCommand):
         sb6 = Submode.objects.create(descriptor="Freight Port", mode=m2)
         sb7 = Submode.objects.create(descriptor="Shipping Infrastructure", mode=m2)
 
+        submodes = [sb1, sb2, sb3, sb4, sb5, sb6, sb7]
+
         # TODO: Create 40 odd organisations here, get their ids and pass them into PersonFactory.create_batch()
         #       below.  Then we need to write a post_generation hook in PersonFactory which ensures that the person
         #       is only added to these Organisations and no further Organisation objects are created.
 
-        number = options["number"][0]
-        user = UserFactory.create()  # we need to have at least one user for the updated_by field
-        role = RoleFactory.create()  # because we have a many-to-many relationship with Role, we need to create one and pass it in
-        PersonFactory.create_batch(number, role=role, updated_by=user,
-                                   predecessor__predecessor=None, organisation__submode=sb1)  # we do this so we don't get a loop
+        # we need a User object to completed the updated_by fields in Organisaton and Person
+        user = (
+            UserFactory.create()
+        )  # we need to have at least one user for the updated_by field
+
+        # Create 40 Organisation objects
+        orgs = [OrganisationFactory.create(submode=submodes[randint(0, len(submodes)-1)]) for org in range(40)]
+
+        role = (
+            RoleFactory.create()
+        )  # because we have a many-to-many relationship with Role, we need to create one and pass it in
+        for org in orgs:
+            PersonFactory.create(
+                role=role,
+                updated_by=user,
+                predecessor=None,
+                organisation__submode=submodes[randint(0, len(submodes)-1)],
+                organisation=org,
+            )
         self.stdout.write(
             self.style.SUCCESS(
                 f"Created {number} Person object[s]! Go forth and multiply."