mirror of https://github.com/zulip/zulip.git
436 lines
20 KiB
Python
436 lines
20 KiB
Python
from __future__ import absolute_import
|
|
from __future__ import division
|
|
from __future__ import print_function
|
|
|
|
from django.core.management.base import BaseCommand, CommandParser
|
|
from django.utils.timezone import now
|
|
|
|
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
|
|
from zerver.lib.actions import STREAM_ASSIGNMENT_COLORS, do_send_messages, \
|
|
do_change_is_admin
|
|
from django.conf import settings
|
|
from zerver.lib.bulk_create import bulk_create_realms, \
|
|
bulk_create_streams, bulk_create_users, bulk_create_huddles, \
|
|
bulk_create_clients
|
|
from zerver.models import DefaultStream, get_stream, get_realm
|
|
from zilencer.models import Deployment
|
|
|
|
import random
|
|
import os
|
|
from optparse import make_option
|
|
from six.moves import range
|
|
from typing import Any, Callable, Dict, List, Iterable, Mapping, Sequence, Set, Tuple, Text
|
|
|
|
settings.TORNADO_SERVER = None
|
|
|
|
def create_users(realm, name_list, bot_type=None):
|
|
# type: (Realm, Iterable[Tuple[Text, Text]], int) -> None
|
|
user_set = set() # type: Set[Tuple[Text, Text, Text, bool]]
|
|
for full_name, email in name_list:
|
|
short_name = email_to_username(email)
|
|
user_set.add((email, full_name, short_name, True))
|
|
tos_version = settings.TOS_VERSION if bot_type is None else None
|
|
bulk_create_users(realm, user_set, bot_type=bot_type, tos_version=tos_version)
|
|
|
|
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.')
|
|
|
|
parser.add_argument('--extra-users',
|
|
dest='extra_users',
|
|
type=int,
|
|
default=0,
|
|
help='The number of extra users to create')
|
|
|
|
parser.add_argument('--extra-bots',
|
|
dest='extra_bots',
|
|
type=int,
|
|
default=0,
|
|
help='The number of extra bots to create')
|
|
|
|
parser.add_argument('--huddles',
|
|
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.')
|
|
|
|
parser.add_argument('--threads',
|
|
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.')
|
|
|
|
parser.add_argument('--percent-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.')
|
|
|
|
parser.add_argument('--nodelete',
|
|
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.')
|
|
|
|
def handle(self, **options):
|
|
# type: (**Any) -> None
|
|
if options["percent_huddles"] + options["percent_personals"] > 100:
|
|
self.stderr.write("Error! More than 100% of messages allocated.\n")
|
|
return
|
|
|
|
if options["delete"]:
|
|
# Start by clearing all the data in our database
|
|
clear_database()
|
|
|
|
# Create our two default realms
|
|
# Could in theory be done via zerver.lib.actions.do_create_realm, but
|
|
# welcome-bot (needed for do_create_realm) hasn't been created yet
|
|
zulip_realm = Realm.objects.create(
|
|
string_id="zulip", name="Zulip Dev", restricted_to_domain=True,
|
|
invite_required=False, org_type=Realm.CORPORATE, domain="zulip.com")
|
|
RealmAlias.objects.create(realm=zulip_realm, domain="zulip.com")
|
|
if options["test_suite"]:
|
|
mit_realm = Realm.objects.create(
|
|
string_id="mit", name="MIT", restricted_to_domain=True,
|
|
invite_required=False, org_type=Realm.CORPORATE, domain="mit.edu")
|
|
RealmAlias.objects.create(realm=mit_realm, domain="mit.edu")
|
|
|
|
# Create test Users (UserProfiles are automatically created,
|
|
# as are subscriptions to the ability to receive personals).
|
|
names = [
|
|
("Zoe", "ZOE@zulip.com"),
|
|
("Othello, the Moor of Venice", "othello@zulip.com"),
|
|
("Iago", "iago@zulip.com"),
|
|
("Prospero from The Tempest", "prospero@zulip.com"),
|
|
("Cordelia Lear", "cordelia@zulip.com"),
|
|
("King Hamlet", "hamlet@zulip.com"),
|
|
("aaron", "AARON@zulip.com"),
|
|
]
|
|
for i in range(options["extra_users"]):
|
|
names.append(('Extra User %d' % (i,), 'extrauser%d@zulip.com' % (i,)))
|
|
create_users(zulip_realm, names)
|
|
iago = UserProfile.objects.get(email="iago@zulip.com")
|
|
do_change_is_admin(iago, True)
|
|
# Create public streams.
|
|
stream_list = ["Verona", "Denmark", "Scotland", "Venice", "Rome"]
|
|
stream_dict = {
|
|
"Verona": {"description": "A city in Italy", "invite_only": False},
|
|
"Denmark": {"description": "A Scandinavian country", "invite_only": False},
|
|
"Scotland": {"description": "Located in the United Kingdom", "invite_only": False},
|
|
"Venice": {"description": "A northeastern Italian city", "invite_only": False},
|
|
"Rome": {"description": "Yet another Italian city", "invite_only": False}
|
|
} # type: Dict[Text, Dict[Text, Any]]
|
|
|
|
bulk_create_streams(zulip_realm, stream_dict)
|
|
recipient_streams = [Stream.objects.get(name=name, realm=zulip_realm).id
|
|
for name in stream_list] # type: List[int]
|
|
# Create subscriptions to streams. The following
|
|
# algorithm will give each of the users a different but
|
|
# deterministic subset of the streams (given a fixed list
|
|
# of users).
|
|
subscriptions_to_add = [] # type: List[Subscription]
|
|
profiles = UserProfile.objects.select_related().all().order_by("email") # type: Sequence[UserProfile]
|
|
for i, profile in enumerate(profiles):
|
|
# Subscribe to some streams.
|
|
for type_id in recipient_streams[:int(len(recipient_streams) *
|
|
float(i)/len(profiles)) + 1]:
|
|
r = Recipient.objects.get(type=Recipient.STREAM, type_id=type_id)
|
|
s = Subscription(
|
|
recipient=r,
|
|
user_profile=profile,
|
|
color=STREAM_ASSIGNMENT_COLORS[i % len(STREAM_ASSIGNMENT_COLORS)])
|
|
|
|
subscriptions_to_add.append(s)
|
|
Subscription.objects.bulk_create(subscriptions_to_add)
|
|
else:
|
|
zulip_realm = get_realm("zulip")
|
|
recipient_streams = [klass.type_id for klass in
|
|
Recipient.objects.filter(type=Recipient.STREAM)]
|
|
|
|
# Extract a list of all users
|
|
user_profiles = list(UserProfile.objects.all()) # type: List[UserProfile]
|
|
|
|
if not options["test_suite"]:
|
|
# Populate users with some bar data
|
|
for user in user_profiles:
|
|
status = UserPresence.ACTIVE # type: int
|
|
date = now()
|
|
client = get_client("website")
|
|
if user.full_name[0] <= 'H':
|
|
client = get_client("ZulipAndroid")
|
|
UserPresence.objects.get_or_create(user_profile=user, client=client, timestamp=date, status=status)
|
|
|
|
user_profiles_ids = [user_profile.id for user_profile in user_profiles]
|
|
|
|
# Create several initial huddles
|
|
for i in range(options["num_huddles"]):
|
|
get_huddle(random.sample(user_profiles_ids, random.randint(3, 4)))
|
|
|
|
# Create several initial pairs for personals
|
|
personals_pairs = [random.sample(user_profiles_ids, 2)
|
|
for i in range(options["num_personals"])]
|
|
|
|
threads = options["threads"]
|
|
jobs = [] # type: List[Tuple[int, List[List[int]], Dict[str, Any], Callable[[str], int]]]
|
|
for i in range(threads):
|
|
count = options["num_messages"] // threads
|
|
if i < options["num_messages"] % threads:
|
|
count += 1
|
|
jobs.append((count, personals_pairs, options, self.stdout.write))
|
|
|
|
for job in jobs:
|
|
send_messages(job)
|
|
|
|
if options["delete"]:
|
|
# Create the "website" and "API" clients; if we don't, the
|
|
# default values in zerver/decorators.py will not work
|
|
# with the Django test suite.
|
|
get_client("website")
|
|
get_client("API")
|
|
|
|
if options["test_suite"]:
|
|
# Create test users; the MIT ones are needed to test
|
|
# the Zephyr mirroring codepaths.
|
|
testsuite_mit_users = [
|
|
("Fred Sipb (MIT)", "sipbtest@mit.edu"),
|
|
("Athena Consulting Exchange User (MIT)", "starnine@mit.edu"),
|
|
("Esp Classroom (MIT)", "espuser@mit.edu"),
|
|
]
|
|
create_users(mit_realm, testsuite_mit_users)
|
|
|
|
# These bots are directly referenced from code and thus
|
|
# are needed for the test suite.
|
|
all_realm_bots = [(bot['name'], bot['email_template'] % (settings.INTERNAL_BOT_DOMAIN,))
|
|
for bot in settings.INTERNAL_BOTS]
|
|
zulip_realm_bots = [
|
|
("Zulip New User Bot", "new-user-bot@zulip.com"),
|
|
("Zulip Error Bot", "error-bot@zulip.com"),
|
|
("Zulip Default Bot", "default-bot@zulip.com"),
|
|
]
|
|
for i in range(options["extra_bots"]):
|
|
zulip_realm_bots.append(('Extra Bot %d' % (i,), 'extrabot%d@zulip.com' % (i,)))
|
|
zulip_realm_bots.extend(all_realm_bots)
|
|
create_users(zulip_realm, zulip_realm_bots, bot_type=UserProfile.DEFAULT_BOT)
|
|
|
|
zulip_webhook_bots = [
|
|
("Zulip Webhook Bot", "webhook-bot@zulip.com"),
|
|
]
|
|
create_users(zulip_realm, zulip_webhook_bots, bot_type=UserProfile.INCOMING_WEBHOOK_BOT)
|
|
|
|
create_simple_community_realm()
|
|
|
|
if not options["test_suite"]:
|
|
# Initialize the email gateway bot as an API Super User
|
|
email_gateway_bot = UserProfile.objects.get(email__iexact=settings.EMAIL_GATEWAY_BOT)
|
|
email_gateway_bot.is_api_super_user = True
|
|
email_gateway_bot.save()
|
|
|
|
# To keep the messages.json fixtures file for the test
|
|
# suite fast, don't add these users and subscriptions
|
|
# when running populate_db for the test suite
|
|
|
|
zulip_stream_dict = {
|
|
"devel": {"description": "For developing", "invite_only": False},
|
|
"all": {"description": "For everything", "invite_only": False},
|
|
"announce": {"description": "For announcements", "invite_only": False},
|
|
"design": {"description": "For design", "invite_only": False},
|
|
"support": {"description": "For support", "invite_only": False},
|
|
"social": {"description": "For socializing", "invite_only": False},
|
|
"test": {"description": "For testing", "invite_only": False},
|
|
"errors": {"description": "For errors", "invite_only": False},
|
|
"sales": {"description": "For sales discussion", "invite_only": False}
|
|
} # type: Dict[Text, Dict[Text, Any]]
|
|
bulk_create_streams(zulip_realm, zulip_stream_dict)
|
|
# Now that we've created the notifications stream, configure it properly.
|
|
zulip_realm.notifications_stream = get_stream("announce", zulip_realm)
|
|
zulip_realm.save(update_fields=['notifications_stream'])
|
|
|
|
# Add a few default streams
|
|
for default_stream_name in ["design", "devel", "social", "support"]:
|
|
DefaultStream.objects.create(realm=zulip_realm,
|
|
stream=get_stream(default_stream_name, zulip_realm))
|
|
|
|
# Now subscribe everyone to these streams
|
|
subscriptions_to_add = []
|
|
profiles = UserProfile.objects.select_related().filter(realm=zulip_realm)
|
|
for i, stream_name in enumerate(zulip_stream_dict):
|
|
stream = Stream.objects.get(name=stream_name, realm=zulip_realm)
|
|
recipient = Recipient.objects.get(type=Recipient.STREAM, type_id=stream.id)
|
|
for profile in profiles:
|
|
# Subscribe to some streams.
|
|
s = Subscription(
|
|
recipient=recipient,
|
|
user_profile=profile,
|
|
color=STREAM_ASSIGNMENT_COLORS[i % len(STREAM_ASSIGNMENT_COLORS)])
|
|
subscriptions_to_add.append(s)
|
|
Subscription.objects.bulk_create(subscriptions_to_add)
|
|
|
|
# These bots are not needed by the test suite
|
|
internal_zulip_users_nosubs = [
|
|
("Zulip Commit Bot", "commit-bot@zulip.com"),
|
|
("Zulip Trac Bot", "trac-bot@zulip.com"),
|
|
("Zulip Nagios Bot", "nagios-bot@zulip.com"),
|
|
]
|
|
create_users(zulip_realm, internal_zulip_users_nosubs, bot_type=UserProfile.DEFAULT_BOT)
|
|
|
|
zulip_cross_realm_bots = [
|
|
("Zulip Feedback Bot", "feedback@zulip.com"),
|
|
]
|
|
create_users(zulip_realm, zulip_cross_realm_bots, bot_type=UserProfile.DEFAULT_BOT)
|
|
|
|
# Mark all messages as read
|
|
UserMessage.objects.all().update(flags=UserMessage.flags.read)
|
|
|
|
self.stdout.write("Successfully populated test database.\n")
|
|
|
|
recipient_hash = {} # type: Dict[int, Recipient]
|
|
def get_recipient_by_id(rid):
|
|
# type: (int) -> Recipient
|
|
if rid in recipient_hash:
|
|
return recipient_hash[rid]
|
|
return Recipient.objects.get(id=rid)
|
|
|
|
# Create some test messages, including:
|
|
# - multiple streams
|
|
# - multiple subjects per stream
|
|
# - multiple huddles
|
|
# - multiple personals converastions
|
|
# - multiple messages per subject
|
|
# - both single and multi-line content
|
|
def send_messages(data):
|
|
# type: (Tuple[int, Sequence[Sequence[int]], Mapping[str, Any], Callable[[str], Any]]) -> int
|
|
(tot_messages, personals_pairs, options, output) = data
|
|
random.seed(os.getpid())
|
|
texts = open("zilencer/management/commands/test_messages.txt", "r").readlines()
|
|
offset = random.randint(0, len(texts))
|
|
|
|
recipient_streams = [klass.id for klass in
|
|
Recipient.objects.filter(type=Recipient.STREAM)] # type: List[int]
|
|
recipient_huddles = [h.id for h in Recipient.objects.filter(type=Recipient.HUDDLE)] # type: List[int]
|
|
|
|
huddle_members = {} # type: Dict[int, List[int]]
|
|
for h in recipient_huddles:
|
|
huddle_members[h] = [s.user_profile.id for s in
|
|
Subscription.objects.filter(recipient_id=h)]
|
|
|
|
num_messages = 0
|
|
random_max = 1000000
|
|
recipients = {} # type: Dict[int, Tuple[int, int, Dict[str, Any]]]
|
|
while num_messages < tot_messages:
|
|
saved_data = {} # type: Dict[str, Any]
|
|
message = Message()
|
|
message.sending_client = get_client('populate_db')
|
|
length = random.randint(1, 5)
|
|
lines = (t.strip() for t in texts[offset: offset + length])
|
|
message.content = '\n'.join(lines)
|
|
offset += length
|
|
offset = offset % len(texts)
|
|
|
|
randkey = random.randint(1, random_max)
|
|
if (num_messages > 0 and
|
|
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:
|
|
personals_pair = saved_data['personals_pair']
|
|
random.shuffle(personals_pair)
|
|
elif message_type == Recipient.STREAM:
|
|
message.subject = saved_data['subject']
|
|
message.recipient = get_recipient_by_id(recipient_id)
|
|
elif message_type == Recipient.HUDDLE:
|
|
message.recipient = get_recipient_by_id(recipient_id)
|
|
elif (randkey <= random_max * options["percent_huddles"] / 100.):
|
|
message_type = Recipient.HUDDLE
|
|
message.recipient = get_recipient_by_id(random.choice(recipient_huddles))
|
|
elif (randkey <= random_max * (options["percent_huddles"] + options["percent_personals"]) / 100.):
|
|
message_type = Recipient.PERSONAL
|
|
personals_pair = random.choice(personals_pairs)
|
|
random.shuffle(personals_pair)
|
|
elif (randkey <= random_max * 1.0):
|
|
message_type = Recipient.STREAM
|
|
message.recipient = get_recipient_by_id(random.choice(recipient_streams))
|
|
|
|
if message_type == Recipient.HUDDLE:
|
|
sender_id = random.choice(huddle_members[message.recipient.id])
|
|
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])
|
|
message.sender = get_user_profile_by_id(personals_pair[1])
|
|
saved_data['personals_pair'] = personals_pair
|
|
elif message_type == Recipient.STREAM:
|
|
stream = Stream.objects.get(id=message.recipient.type_id)
|
|
# Pick a random subscriber to the stream
|
|
message.sender = random.choice(Subscription.objects.filter(
|
|
recipient=message.recipient)).user_profile
|
|
message.subject = stream.name + Text(random.randint(1, 3))
|
|
saved_data['subject'] = message.subject
|
|
|
|
message.pub_date = now()
|
|
do_send_messages([{'message': message}])
|
|
|
|
recipients[num_messages] = (message_type, message.recipient.id, saved_data)
|
|
num_messages += 1
|
|
return tot_messages
|
|
|
|
def create_simple_community_realm():
|
|
# type: () -> None
|
|
simple_realm = Realm.objects.create(
|
|
string_id="simple", name="Simple Realm", restricted_to_domain=False,
|
|
invite_required=False, org_type=Realm.COMMUNITY, domain="simple.com")
|
|
|
|
names = [
|
|
("alice", "alice@example.com"),
|
|
("bob", "bob@foo.edu"),
|
|
("cindy", "cindy@foo.tv"),
|
|
]
|
|
create_users(simple_realm, names)
|
|
|
|
user_profiles = UserProfile.objects.filter(realm__string_id='simple')
|
|
create_user_presences(user_profiles)
|
|
|
|
def create_user_presences(user_profiles):
|
|
# type: (Iterable[UserProfile]) -> None
|
|
for user in user_profiles:
|
|
status = 1 # type: int
|
|
date = now()
|
|
client = get_client("website")
|
|
UserPresence.objects.get_or_create(
|
|
user_profile=user,
|
|
client=client,
|
|
timestamp=date,
|
|
status=status)
|