Updated reset data script for locust

dlareau · May 3, 2020 · 25f4eef · 25f4eef
1 parent c2da438
commit 25f4eef
Showing 1 changed file with 24 additions and 9 deletions.
diff --git a/locust/reset_data.py b/locust/reset_data.py
@@ -1,8 +1,11 @@
 import os
 import sys
 
-sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
-os.environ.setdefault("DJANGO_SETTINGS_MODULE", "puzzlehunt_server.settings.local_settings")
+# sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+os.environ.setdefault("DJANGO_SETTINGS_MODULE", "puzzlehunt_server.settings.env_settings")
+
+import django
+django.setup()
 
 # your imports, e.g. Django models
 from django.contrib.auth.models import User
@@ -19,7 +22,7 @@
     if not ans:
         sys.exit()
     if ans not in ['y', 'Y', 'n', 'N']:
-        print 'please enter y or n.'
+        print('please enter y or n.')
         continue
     if ans == 'n' or ans == 'N':
         sys.exit()
@@ -28,24 +31,36 @@
 
 # Wipe existing data and load base data
 call_command('flush', verbosity=0, interactive=False)
-call_command('loaddata', 'locust', verbosity=0)
+call_command('loaddata', 'locust', ignorenonexistent=True)
 
 # Add in the requested number of users/people/teams
 curr_hunt = Hunt.objects.get(is_current_hunt=True)
 
 team_size = curr_hunt.team_size
-num_bots = 800  # should be a multiple of 25
-num_staff = num_bots / 25
+num_bots = 900  # should be a multiple of 25 and of team size
+num_staff = int(num_bots / 25)
 num_players = num_bots - num_staff
-num_teams = num_players % team_size
+num_teams = int(num_players / team_size)
+
+print("Now creating %d players, %d teams, and %d staff members" % (num_players, num_teams, num_staff))
 
 for team_index in range(num_teams):
+    sys.stdout.write(".")
     team = Team.objects.create(team_name="team_" + str(team_index), hunt=curr_hunt,
                                location="test_loc", join_code="JOIN1")
     for person_index in range(team_size):
         user_number = team_index * team_size + person_index
         user = User.objects.create_user('test_user_' + str(user_number),
                                         'example' + str(user_number) + '@example.com',
                                         'password' + str(user_number))
-        Person.objects.create(user=user, is_shib_acct=False)
-        Person.teams.add(team)
+        person = Person.objects.create(user=user, is_shib_acct=False)
+        person.teams.add(team)
+        person.save()
+for person_index in range(num_staff):
+    user_number = person_index + num_players
+    user = User.objects.create_superuser('test_user_' + str(user_number),
+                                         'example' + str(user_number) + '@example.com',
+                                         'password' + str(user_number))
+    person = Person.objects.create(user=user, is_shib_acct=False)
+    person.save()
+print("\n")