From 9efdb1e77c623e5d6d25cfa12f2fe83c92cf9390 Mon Sep 17 00:00:00 2001 From: Juan Verhook Date: Wed, 7 Dec 2016 01:13:25 +0000 Subject: [PATCH] Added UserPresence to population script --- zilencer/management/commands/populate_db.py | 104 +++++++++++--------- 1 file changed, 60 insertions(+), 44 deletions(-) diff --git a/zilencer/management/commands/populate_db.py b/zilencer/management/commands/populate_db.py index 300cbeb436d54e..ac6e34fd83f0a4 100644 --- a/zilencer/management/commands/populate_db.py +++ b/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: