improved populate script now

2 files changed, 35 insertions, 21 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."
diff --git a/ctrack/organisations/tests/factories.py b/ctrack/organisations/tests/factories.py
index 3bad765..61f2ae3 100644
--- a/ctrack/organisations/tests/factories.py
+++ b/ctrack/organisations/tests/factories.py
@@ -40,28 +40,13 @@ class UserFactory(DjangoModelFactory):
     password = Faker("lexify", text="????????", letters="AdOIqkcvBnMP")
 
 
-class ModeFactory(DjangoModelFactory):
-    class Meta:
-        model = Mode
-
-    descriptor = factory.LazyFunction(_random_mode)
-
-
-class SubModeFactory(DjangoModelFactory):
-    class Meta:
-        model = Submode
-
-    descriptor = factory.LazyFunction(_random_submode)
-    mode = SubFactory(ModeFactory)
-
-
 class OrganisationFactory(DjangoModelFactory):
     class Meta:
         model = Organisation
 
     name = Faker("company")
     slug = Faker("lexify", text="????", letters="abcdsg")
-    submode = SubFactory(SubModeFactory)
+#   submode = SubFactory(SubModeFactory)
     designation_type = 1
     registered_company_name = Faker("company")
     registered_company_number = Faker("numerify", text="######")
@@ -82,6 +67,14 @@ class PersonFactory(DjangoModelFactory):
     class Meta:
         model = Person
 
+    # @post_generation
+    # def organisation(self, create, extracted, **kwargs):
+    #     if not create:
+    #         return
+    #     if extracted:
+    #         for org in extracted:
+    #             self.organisation.add(org)
+
     @post_generation
     def role(self, create, extracted, **kwargs):
         if not create: