From 527058a202510b00a481aa2172a5f6017fcc9fd7 Mon Sep 17 00:00:00 2001 From: Tim Abbott Date: Wed, 5 Sep 2012 12:38:35 -0400 Subject: [PATCH] Add personals and huddles to populate_db. (imported from commit 98a8decc18bad4700260f125873ccc4367f73ebe) --- zephyr/management/commands/populate_db.py | 112 ++++++++++++++++------ 1 file changed, 84 insertions(+), 28 deletions(-) diff --git a/zephyr/management/commands/populate_db.py b/zephyr/management/commands/populate_db.py index d0f2678178..951e3e02af 100644 --- a/zephyr/management/commands/populate_db.py +++ b/zephyr/management/commands/populate_db.py @@ -2,7 +2,8 @@ from django.core.management.base import BaseCommand from django.utils.timezone import utc from django.contrib.auth.models import User -from zephyr.models import Zephyr, UserProfile, ZephyrClass, Recipient, Subscription +from zephyr.models import Zephyr, UserProfile, ZephyrClass, Recipient, \ + Subscription, Huddle, get_huddle import datetime import random @@ -15,12 +16,37 @@ class Command(BaseCommand): make_option('-n', '--num-zephyrs', dest='num_zephyrs', type='int', - default=100, + default=120, help='The number of zephyrs to create.'), + make_option('--huddles', + dest='num_huddles', + type='int', + default=3, + help='The number of huddles to create.'), + make_option('--personals', + dest='num_personals', + type='int', + default=6, + help='The number of personal pairs to create.'), + make_option('--percent-huddles', + dest='percent_huddles', + type='float', + default=15, + help='The percent of messages to be huddles.'), + make_option('--percent-personals', + dest='percent_personals', + type='float', + default=15, + help='The percent of messages to be personals.'), + ) def handle(self, **options): - for klass in [Zephyr, ZephyrClass, UserProfile, User, Recipient, Subscription]: + if options["percent_huddles"] + options["percent_personals"] > 100: + self.stderr.write("Error! More than 100% of messages allocated.\n") + return + + for klass in [Zephyr, ZephyrClass, UserProfile, User, Recipient, Subscription, Huddle]: klass.objects.all().delete() # Create test Users (UserProfiles are automatically created, @@ -29,6 +55,7 @@ class Command(BaseCommand): for username in usernames: u = User.objects.create_user(username=username, password=username) u.save() + users = [user.id for user in User.objects.all()] # Create public classes. for name in ["Verona", "Denmark", "Scotland", "Venice", "Rome"]: @@ -38,37 +65,66 @@ class Command(BaseCommand): recipient = Recipient(type_id=new_class.pk, type="class") recipient.save() - # Create some test zephyrs, including: - # - multiple classes - # - multiple instances per class - # - multiple zephyrs per instance - # - both single and multi-line content - users = [user.id for user in User.objects.all()] - recipient_classes = [klass.id for klass in Recipient.objects.filter(type="class")] - texts = file("zephyr/management/commands/test_zephyrs.txt", "r").readlines() - offset = 0 - num_zephyrs = 0 - while num_zephyrs < options["num_zephyrs"]: - new_zephyr = Zephyr() - new_zephyr.sender = UserProfile.objects.get(id=random.choice(users)) - length = random.randint(1, 5) - new_zephyr.content = "".join(texts[offset: offset + length]) - offset += length - offset = offset % len(texts) - new_zephyr.recipient = Recipient.objects.get(id=random.choice(recipient_classes)) - zephyr_class = ZephyrClass.objects.get(pk=new_zephyr.recipient.type_id) - new_zephyr.instance = zephyr_class.name + str(random.randint(1, 3)) - new_zephyr.pub_date = datetime.datetime.utcnow().replace(tzinfo=utc) - new_zephyr.save() - num_zephyrs += 1 + # Create several initial huddles + huddle_members = {} + for i in range(0, options["num_huddles"]): + user_ids = random.sample(users, random.randint(3, 4)) + huddle_members[get_huddle(user_ids).id] = user_ids - # Create subscriptions + # Create several initial pairs for personals + personals_pairs = [] + for i in range(0, options["num_personals"]): + personals_pairs.append(random.sample(users, 2)) + + recipient_classes = [klass.type_id for klass in Recipient.objects.filter(type="class")] + recipient_huddles = [h.type_id for h in Recipient.objects.filter(type="huddle")] + + # Create subscriptions to classes profiles = UserProfile.objects.all() for i, profile in enumerate(profiles): # Subscribe to some classes. for recipient in recipient_classes[:int(len(recipient_classes) * float(i)/len(profiles)) + 1]: new_subscription = Subscription(userprofile_id=profile, - recipient_id=Recipient.objects.get(id=recipient)) + recipient_id=Recipient.objects.get(type="class", + type_id=recipient)) new_subscription.save() + # Create some test zephyrs, including: + # - multiple classes + # - multiple instances per class + # - multiple zephyrs per instance + # - both single and multi-line content + + texts = file("zephyr/management/commands/test_zephyrs.txt", "r").readlines() + offset = 0 + num_zephyrs = 0 + random_max = 1000000 + while num_zephyrs < options["num_zephyrs"]: + new_zephyr = Zephyr() + length = random.randint(1, 5) + new_zephyr.content = "".join(texts[offset: offset + length]) + offset += length + offset = offset % len(texts) + + randkey = random.randint(1, random_max) + if (randkey <= random_max * options["percent_huddles"] / 100.): + # huddle + new_zephyr.recipient = Recipient.objects.get(type="huddle", type_id=random.choice(recipient_huddles)) + new_zephyr.sender = UserProfile.objects.get(id=random.choice(huddle_members[new_zephyr.recipient.type_id])) + elif (randkey <= random_max * (options["percent_huddles"] + options["percent_personals"]) / 100.): + # personals + pair = random.choice(personals_pairs) + random.shuffle(pair) + new_zephyr.recipient = Recipient.objects.get(type="personal", type_id=pair[0]) + new_zephyr.sender = UserProfile.objects.get(id=pair[1]) + elif (randkey <= random_max * 1.0): + # class + new_zephyr.recipient = Recipient.objects.get(type="class", type_id=random.choice(recipient_classes)) + zephyr_class = ZephyrClass.objects.get(pk=new_zephyr.recipient.type_id) + new_zephyr.sender = UserProfile.objects.get(id=random.choice(users)) + new_zephyr.instance = zephyr_class.name + str(random.randint(1, 3)) + new_zephyr.pub_date = datetime.datetime.utcnow().replace(tzinfo=utc) + new_zephyr.save() + num_zephyrs += 1 + self.stdout.write("Successfully populated test database.\n")