Added UserPresence to population script

zilencer/management/commands/populate_db.py

@@ -5,7 +5,7 @@
 from django.core.management.base import BaseCommand, CommandParser
 from django.utils.timezone import now
 
-from zerver.models import Message, UserProfile, Stream, Recipient, \
+from zerver.models import Message, UserProfile, Stream, Recipient, UserPresence, \
     Subscription, get_huddle, Realm, UserMessage, RealmAlias, \
     clear_database, get_client, get_user_profile_by_id, \
     email_to_username
@@ -15,6 +15,7 @@
 from zerver.lib.bulk_create import bulk_create_streams, bulk_create_users
 from zerver.models import DefaultStream, get_stream, get_realm
 from zilencer.models import Deployment
+import datetime
 
 import random
 import os
@@ -41,69 +42,70 @@ def create_streams(realms, realm, stream_list):
         stream_set.add((realm.domain, stream_name))
     bulk_create_streams(realms, stream_set)
 
+
 class Command(BaseCommand):
     help = "Populate a test database"
 
     def add_arguments(self, parser):
         # type: (CommandParser) -> None
         parser.add_argument('-n', '--num-messages',
-                    dest='num_messages',
-                    type=int,
-                    default=600,
-                    help='The number of messages to create.')
+                            dest='num_messages',
+                            type=int,
+                            default=600,
+                            help='The number of messages to create.')
 
         parser.add_argument('--extra-users',
-                    dest='extra_users',
-                    type=int,
-                    default=0,
-                    help='The number of extra users to create')
+                            dest='extra_users',
+                            type=int,
+                            default=0,
+                            help='The number of extra users to create')
 
         parser.add_argument('--huddles',
-                    dest='num_huddles',
-                    type=int,
-                    default=3,
-                    help='The number of huddles to create.')
+                            dest='num_huddles',
+                            type=int,
+                            default=3,
+                            help='The number of huddles to create.')
 
         parser.add_argument('--personals',
-                    dest='num_personals',
-                    type=int,
-                    default=6,
-                    help='The number of personal pairs to create.')
+                            dest='num_personals',
+                            type=int,
+                            default=6,
+                            help='The number of personal pairs to create.')
 
         parser.add_argument('--threads',
-                    dest='threads',
-                    type=int,
-                    default=10,
-                    help='The number of threads to use.')
+                            dest='threads',
+                            type=int,
+                            default=10,
+                            help='The number of threads to use.')
 
         parser.add_argument('--percent-huddles',
-                    dest='percent_huddles',
-                    type=float,
-                    default=15,
-                    help='The percent of messages to be huddles.')
+                            dest='percent_huddles',
+                            type=float,
+                            default=15,
+                            help='The percent of messages to be huddles.')
 
         parser.add_argument('--percent-personals',
-                    dest='percent_personals',
-                    type=float,
-                    default=15,
-                    help='The percent of messages to be personals.')
+                            dest='percent_personals',
+                            type=float,
+                            default=15,
+                            help='The percent of messages to be personals.')
 
         parser.add_argument('--stickyness',
-                    dest='stickyness',
-                    type=float,
-                    default=20,
-                    help='The percent of messages to repeat recent folks.')
+                            dest='stickyness',
+                            type=float,
+                            default=20,
+                            help='The percent of messages to repeat recent folks.')
 
         parser.add_argument('--nodelete',
-                    action="store_false",
-                    default=True,
-                    dest='delete',
-                    help='Whether to delete all the existing messages.')
+                            action="store_false",
+                            default=True,
+                            dest='delete',
+                            help='Whether to delete all the existing messages.')
 
         parser.add_argument('--test-suite',
-                    default=False,
-                    action="store_true",
-                    help='Whether to delete all the existing messages.')
+                            default=False,
+                            action="store_true",
+                            help='Whether to delete all the existing messages.')
 
     def handle(self, **options):
         # type: (**Any) -> None
@@ -177,8 +179,22 @@ def handle(self, **options):
                                  Recipient.objects.filter(type=Recipient.STREAM)]
 
         # Extract a list of all users
-        user_profiles = [user_profile.id for user_profile in UserProfile.objects.all()] # type: List[int]
-
+        user_profiles = [user_profile for user_profile in UserProfile.objects.all()] # type: List[Any]
+
+        # Populate users with some bar data
+        for user in user_profiles:
+            month = random.choice(range(1, 13)) # type: Any
+            day = random.choice(range(1, 29)) # type: Any
+            random_status = 2 # type: int
+            random_date = datetime.datetime(2015, month, day) # type: datetime.datetime
+            random_client = random.choice(range(1, 3)) # type: Any
+            if random_client == 1:
+                random_client = get_client("API")
+            else:
+                random_client = get_client("website")
+            UserPresence.objects.get_or_create(user_profile=user, client=random_client, timestamp=random_date, status=random_status)
+
+        user_profiles = [user_profile.id for user_profile in UserProfile.objects.all()]
         # Create several initial huddles
         for i in range(options["num_huddles"]):
             get_huddle(random.sample(user_profiles, random.randint(3, 4)))
@@ -328,7 +344,7 @@ def send_messages(data):
 
         randkey = random.randint(1, random_max)
         if (num_messages > 0 and
-            random.randint(1, random_max) * 100. / random_max < options["stickyness"]):
+                random.randint(1, random_max) * 100. / random_max < options["stickyness"]):
             # Use an old recipient
             message_type, recipient_id, saved_data = recipients[num_messages - 1]
             if message_type == Recipient.PERSONAL:
@@ -355,7 +371,7 @@ def send_messages(data):
             message.sender = get_user_profile_by_id(sender_id)
         elif message_type == Recipient.PERSONAL:
             message.recipient = Recipient.objects.get(type=Recipient.PERSONAL,
-                                                         type_id=personals_pair[0])
+                                                      type_id=personals_pair[0])
             message.sender = get_user_profile_by_id(personals_pair[1])
             saved_data['personals_pair'] = personals_pair
         elif message_type == Recipient.STREAM: