Example #1
0
    def handle(self, **options: Any) -> None:
        if options["percent_huddles"] + options["percent_personals"] > 100:
            self.stderr.write(
                "Error!  More than 100% of messages allocated.\n")
            return

        # Get consistent data for backend tests.
        if options["test_suite"]:
            random.seed(0)

        if options["delete"]:
            # Start by clearing all the data in our database
            clear_database()

            # Create our three 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
            create_internal_realm()
            zulip_realm = Realm.objects.create(
                string_id="zulip",
                name="Zulip Dev",
                emails_restricted_to_domains=True,
                description=
                "The Zulip development environment default organization."
                "  It's great for testing!",
                invite_required=False,
                org_type=Realm.CORPORATE)
            RealmDomain.objects.create(realm=zulip_realm, domain="zulip.com")
            if options["test_suite"]:
                mit_realm = Realm.objects.create(
                    string_id="zephyr",
                    name="MIT",
                    emails_restricted_to_domains=True,
                    invite_required=False,
                    org_type=Realm.CORPORATE)
                RealmDomain.objects.create(realm=mit_realm, domain="mit.edu")

                lear_realm = Realm.objects.create(
                    string_id="lear",
                    name="Lear & Co.",
                    emails_restricted_to_domains=False,
                    invite_required=False,
                    org_type=Realm.CORPORATE)

            # Create test Users (UserProfiles are automatically created,
            # as are subscriptions to the ability to receive personals).
            names = [
                ("Zoe", "*****@*****.**"),
                ("Othello, the Moor of Venice", "*****@*****.**"),
                ("Iago", "*****@*****.**"),
                ("Prospero from The Tempest", "*****@*****.**"),
                ("Cordelia Lear", "*****@*****.**"),
                ("King Hamlet", "*****@*****.**"),
                ("aaron", "*****@*****.**"),
                ("Polonius", "*****@*****.**"),
            ]

            # For testing really large batches:
            # Create extra users with semi realistic names to make search
            # functions somewhat realistic.  We'll still create 1000 users
            # like Extra222 User for some predicability.
            num_names = options['extra_users']
            num_boring_names = 1000

            for i in range(min(num_names, num_boring_names)):
                full_name = 'Extra%03d User' % (i, )
                names.append((full_name, '*****@*****.**' % (i, )))

            if num_names > num_boring_names:
                fnames = [
                    'Amber', 'Arpita', 'Bob', 'Cindy', 'Daniela', 'Dan',
                    'Dinesh', 'Faye', 'François', 'George', 'Hank', 'Irene',
                    'James', 'Janice', 'Jenny', 'Jill', 'John', 'Kate',
                    'Katelyn', 'Kobe', 'Lexi', 'Manish', 'Mark', 'Matt',
                    'Mayna', 'Michael', 'Pete', 'Peter', 'Phil', 'Phillipa',
                    'Preston', 'Sally', 'Scott', 'Sandra', 'Steve',
                    'Stephanie', 'Vera'
                ]
                mnames = ['de', 'van', 'von', 'Shaw', 'T.']
                lnames = [
                    'Adams', 'Agarwal', 'Beal', 'Benson', 'Bonita', 'Davis',
                    'George', 'Harden', 'James', 'Jones', 'Johnson', 'Jordan',
                    'Lee', 'Leonard', 'Singh', 'Smith', 'Patel', 'Towns',
                    'Wall'
                ]

            for i in range(num_boring_names, num_names):
                fname = random.choice(fnames) + str(i)
                full_name = fname
                if random.random() < 0.7:
                    if random.random() < 0.5:
                        full_name += ' ' + random.choice(mnames)
                    full_name += ' ' + random.choice(lnames)
                email = fname.lower() + '@zulip.com'
                names.append((full_name, email))

            create_users(zulip_realm, names, tos_version=settings.TOS_VERSION)

            iago = get_user("*****@*****.**", zulip_realm)
            do_change_is_admin(iago, True)
            iago.is_staff = True
            iago.save(update_fields=['is_staff'])

            guest_user = get_user("*****@*****.**", zulip_realm)
            guest_user.role = UserProfile.ROLE_GUEST
            guest_user.save(update_fields=['role'])

            # These bots are directly referenced from code and thus
            # are needed for the test suite.
            zulip_realm_bots = [
                ("Zulip Error Bot", "*****@*****.**"),
                ("Zulip Default Bot", "*****@*****.**"),
            ]
            for i in range(options["extra_bots"]):
                zulip_realm_bots.append(
                    ('Extra Bot %d' % (i, ), '*****@*****.**' % (i, )))

            create_users(zulip_realm,
                         zulip_realm_bots,
                         bot_type=UserProfile.DEFAULT_BOT)

            zoe = get_user("*****@*****.**", zulip_realm)
            zulip_webhook_bots = [
                ("Zulip Webhook Bot", "*****@*****.**"),
            ]
            # If a stream is not supplied in the webhook URL, the webhook
            # will (in some cases) send the notification as a PM to the
            # owner of the webhook bot, so bot_owner can't be None
            create_users(zulip_realm,
                         zulip_webhook_bots,
                         bot_type=UserProfile.INCOMING_WEBHOOK_BOT,
                         bot_owner=zoe)
            aaron = get_user("*****@*****.**", zulip_realm)

            zulip_outgoing_bots = [("Outgoing Webhook",
                                    "*****@*****.**")]
            create_users(zulip_realm,
                         zulip_outgoing_bots,
                         bot_type=UserProfile.OUTGOING_WEBHOOK_BOT,
                         bot_owner=aaron)
            outgoing_webhook = get_user("*****@*****.**",
                                        zulip_realm)
            add_service("outgoing-webhook",
                        user_profile=outgoing_webhook,
                        interface=Service.GENERIC,
                        base_url="http://127.0.0.1:5002",
                        token=generate_api_key())

            # Add the realm internl bots to each realm.
            create_if_missing_realm_internal_bots()

            # Create public streams.
            stream_list = ["Verona", "Denmark", "Scotland", "Venice", "Rome"]
            stream_dict = {
                "Verona": {
                    "description": "A city in Italy"
                },
                "Denmark": {
                    "description": "A Scandinavian country"
                },
                "Scotland": {
                    "description": "Located in the United Kingdom"
                },
                "Venice": {
                    "description": "A northeastern Italian city"
                },
                "Rome": {
                    "description": "Yet another Italian city",
                    "is_web_public": True
                }
            }  # type: Dict[str, Dict[str, 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). For the test suite, we have a fixed list of
            # subscriptions to make sure test data is consistent
            # across platforms.

            subscriptions_list = [
            ]  # type: List[Tuple[UserProfile, Recipient]]
            profiles = UserProfile.objects.select_related().filter(
                is_bot=False).order_by("email")  # type: Sequence[UserProfile]

            if options["test_suite"]:
                subscriptions_map = {
                    '*****@*****.**': ['Verona'],
                    '*****@*****.**': ['Verona'],
                    '*****@*****.**': ['Verona', 'Denmark'],
                    '*****@*****.**': ['Verona', 'Denmark', 'Scotland'],
                    '*****@*****.**': ['Verona', 'Denmark', 'Scotland'],
                    '*****@*****.**':
                    ['Verona', 'Denmark', 'Scotland', 'Venice'],
                    '*****@*****.**':
                    ['Verona', 'Denmark', 'Scotland', 'Venice', 'Rome'],
                    '*****@*****.**': ['Verona'],
                }

                for profile in profiles:
                    if profile.email not in subscriptions_map:
                        raise Exception(
                            'Subscriptions not listed for user %s' %
                            (profile.email, ))

                    for stream_name in subscriptions_map[profile.email]:
                        stream = Stream.objects.get(name=stream_name)
                        r = Recipient.objects.get(type=Recipient.STREAM,
                                                  type_id=stream.id)
                        subscriptions_list.append((profile, r))
            else:
                num_streams = len(recipient_streams)
                num_users = len(profiles)
                for i, profile in enumerate(profiles):
                    # Subscribe to some streams.
                    fraction = float(i) / num_users
                    num_recips = int(num_streams * fraction) + 1

                    for type_id in recipient_streams[:num_recips]:
                        r = Recipient.objects.get(type=Recipient.STREAM,
                                                  type_id=type_id)
                        subscriptions_list.append((profile, r))

            subscriptions_to_add = []  # type: List[Subscription]
            event_time = timezone_now()
            all_subscription_logs = []  # type: (List[RealmAuditLog])

            i = 0
            for profile, recipient in subscriptions_list:
                i += 1
                color = STREAM_ASSIGNMENT_COLORS[i %
                                                 len(STREAM_ASSIGNMENT_COLORS)]
                s = Subscription(recipient=recipient,
                                 user_profile=profile,
                                 color=color)

                subscriptions_to_add.append(s)

                log = RealmAuditLog(
                    realm=profile.realm,
                    modified_user=profile,
                    modified_stream_id=recipient.type_id,
                    event_last_message_id=0,
                    event_type=RealmAuditLog.SUBSCRIPTION_CREATED,
                    event_time=event_time)
                all_subscription_logs.append(log)

            Subscription.objects.bulk_create(subscriptions_to_add)
            RealmAuditLog.objects.bulk_create(all_subscription_logs)

            # Create custom profile field data
            phone_number = try_add_realm_custom_profile_field(
                zulip_realm,
                "Phone number",
                CustomProfileField.SHORT_TEXT,
                hint='')
            biography = try_add_realm_custom_profile_field(
                zulip_realm,
                "Biography",
                CustomProfileField.LONG_TEXT,
                hint='What are you known for?')
            favorite_food = try_add_realm_custom_profile_field(
                zulip_realm,
                "Favorite food",
                CustomProfileField.SHORT_TEXT,
                hint="Or drink, if you'd prefer")
            field_data = {
                'vim': {
                    'text': 'Vim',
                    'order': '1'
                },
                'emacs': {
                    'text': 'Emacs',
                    'order': '2'
                },
            }  # type: ProfileFieldData
            favorite_editor = try_add_realm_custom_profile_field(
                zulip_realm,
                "Favorite editor",
                CustomProfileField.CHOICE,
                field_data=field_data)
            birthday = try_add_realm_custom_profile_field(
                zulip_realm, "Birthday", CustomProfileField.DATE)
            favorite_website = try_add_realm_custom_profile_field(
                zulip_realm,
                "Favorite website",
                CustomProfileField.URL,
                hint="Or your personal blog's URL")
            mentor = try_add_realm_custom_profile_field(
                zulip_realm, "Mentor", CustomProfileField.USER)
            github_profile = try_add_realm_default_custom_profile_field(
                zulip_realm, "github")

            # Fill in values for Iago and Hamlet
            hamlet = get_user("*****@*****.**", zulip_realm)
            do_update_user_custom_profile_data_if_changed(
                iago, [
                    {
                        "id": phone_number.id,
                        "value": "+1-234-567-8901"
                    },
                    {
                        "id": biography.id,
                        "value": "Betrayer of Othello."
                    },
                    {
                        "id": favorite_food.id,
                        "value": "Apples"
                    },
                    {
                        "id": favorite_editor.id,
                        "value": "emacs"
                    },
                    {
                        "id": birthday.id,
                        "value": "2000-1-1"
                    },
                    {
                        "id": favorite_website.id,
                        "value": "https://zulip.readthedocs.io/en/latest/"
                    },
                    {
                        "id": mentor.id,
                        "value": [hamlet.id]
                    },
                    {
                        "id": github_profile.id,
                        "value": 'zulip'
                    },
                ])
            do_update_user_custom_profile_data_if_changed(
                hamlet, [
                    {
                        "id": phone_number.id,
                        "value": "+0-11-23-456-7890"
                    },
                    {
                        "id":
                        biography.id,
                        "value":
                        "I am:\n* The prince of Denmark\n* Nephew to the usurping Claudius",
                    },
                    {
                        "id": favorite_food.id,
                        "value": "Dark chocolate"
                    },
                    {
                        "id": favorite_editor.id,
                        "value": "vim"
                    },
                    {
                        "id": birthday.id,
                        "value": "1900-1-1"
                    },
                    {
                        "id": favorite_website.id,
                        "value": "https://blog.zulig.org"
                    },
                    {
                        "id": mentor.id,
                        "value": [iago.id]
                    },
                    {
                        "id": github_profile.id,
                        "value": 'zulipbot'
                    },
                ])
        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.filter(
            is_bot=False))  # type: List[UserProfile]

        # Create a test realm emoji.
        IMAGE_FILE_PATH = static_path('images/test-images/checkbox.png')
        with open(IMAGE_FILE_PATH, 'rb') as fp:
            check_add_realm_emoji(zulip_realm, 'green_tick', iago, fp)

        if not options["test_suite"]:
            # Populate users with some bar data
            for user in user_profiles:
                status = UserPresence.ACTIVE  # type: int
                date = timezone_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"])
        ]

        # Generate a new set of test data.
        create_test_data()

        # prepopulate the URL preview/embed data for the links present
        # in the config.generate_data.json data set.  This makes it
        # possible for populate_db to run happily without Internet
        # access.
        with open("zerver/tests/fixtures/docs_url_preview_data.json",
                  "r") as f:
            urls_with_preview_data = ujson.load(f)
            for url in urls_with_preview_data:
                cache_set(url, urls_with_preview_data[url], PREVIEW_CACHE_NAME)

        threads = options["threads"]
        jobs = [
        ]  # type: List[Tuple[int, List[List[int]], Dict[str, Any], Callable[[str], int], 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,
                         random.randint(0, 10**10)))

        for job in jobs:
            generate_and_send_messages(job)

        if options["delete"]:
            if options["test_suite"]:
                # Create test users; the MIT ones are needed to test
                # the Zephyr mirroring codepaths.
                testsuite_mit_users = [
                    ("Fred Sipb (MIT)", "*****@*****.**"),
                    ("Athena Consulting Exchange User (MIT)",
                     "*****@*****.**"),
                    ("Esp Classroom (MIT)", "*****@*****.**"),
                ]
                create_users(mit_realm,
                             testsuite_mit_users,
                             tos_version=settings.TOS_VERSION)

                testsuite_lear_users = [
                    ("King Lear", "*****@*****.**"),
                    ("Cordelia Lear", "*****@*****.**"),
                ]
                create_users(lear_realm,
                             testsuite_lear_users,
                             tos_version=settings.TOS_VERSION)

            if not options["test_suite"]:
                # 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"
                    },
                    "all": {
                        "description": "For **everything**"
                    },
                    "announce": {
                        "description": "For announcements",
                        'stream_post_policy': Stream.STREAM_POST_POLICY_ADMINS
                    },
                    "design": {
                        "description": "For design"
                    },
                    "support": {
                        "description": "For support"
                    },
                    "social": {
                        "description": "For socializing"
                    },
                    "test": {
                        "description": "For testing `code`"
                    },
                    "errors": {
                        "description": "For errors"
                    },
                    "sales": {
                        "description": "For sales discussion"
                    }
                }  # type: Dict[str, Dict[str, Any]]

                # Calculate the maximum number of digits in any extra stream's
                # number, since a stream with name "Extra Stream 3" could show
                # up after "Extra Stream 29". (Used later to pad numbers with
                # 0s).
                maximum_digits = len(str(options['extra_streams'] - 1))

                for i in range(options['extra_streams']):
                    # Pad the number with 0s based on `maximum_digits`.
                    number_str = str(i).zfill(maximum_digits)

                    extra_stream_name = 'Extra Stream ' + number_str

                    zulip_stream_dict[extra_stream_name] = {
                        "description": "Auto-generated extra stream.",
                    }

                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
                subscribe_users_to_streams(zulip_realm, zulip_stream_dict)

                # These bots are not needed by the test suite
                internal_zulip_users_nosubs = [
                    ("Zulip Commit Bot", "*****@*****.**"),
                    ("Zulip Trac Bot", "*****@*****.**"),
                    ("Zulip Nagios Bot", "*****@*****.**"),
                ]
                create_users(zulip_realm,
                             internal_zulip_users_nosubs,
                             bot_type=UserProfile.DEFAULT_BOT)

            # Mark all messages as read
            UserMessage.objects.all().update(flags=UserMessage.flags.read)

            if not options["test_suite"]:
                # Update pointer of each user to point to the last message in their
                # UserMessage rows with sender_id=user_profile_id.
                users = list(
                    UserMessage.objects.filter(message__sender_id=F(
                        'user_profile_id')).values('user_profile_id').annotate(
                            pointer=Max('message_id')))
                for user in users:
                    UserProfile.objects.filter(
                        id=user['user_profile_id']).update(
                            pointer=user['pointer'])

            create_user_groups()

            if not options["test_suite"]:
                # We populate the analytics database here for
                # development purpose only
                call_command('populate_analytics_db')
            self.stdout.write("Successfully populated test database.\n")
Example #2
0
    def handle(self, **options: Any) -> None:
        if options["percent_huddles"] + options["percent_personals"] > 100:
            self.stderr.write(
                "Error!  More than 100% of messages allocated.\n")
            return

        # Get consistent data for backend tests.
        if options["test_suite"]:
            random.seed(0)

            with connection.cursor() as cursor:
                # Sometimes bugs relating to confusing recipient.id for recipient.type_id
                # or <object>.id for <object>.recipient_id remain undiscovered by the test suite
                # due to these numbers happening to coincide in such a way that it makes tests
                # accidentally pass. By bumping the Recipient.id sequence by a large enough number,
                # we can have those ids in a completely different range of values than object ids,
                # eliminatng the possibility of such coincidences.
                cursor.execute("SELECT setval('zerver_recipient_id_seq', 100)")

        # If max_topics is not set, we set it proportional to the
        # number of messages.
        if options["max_topics"] is None:
            options["max_topics"] = 1 + options["num_messages"] // 100

        if options["delete"]:
            # Start by clearing all the data in our database
            clear_database()

            # Create our three 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
            create_internal_realm()
            zulip_realm = do_create_realm(
                string_id="zulip",
                name="Zulip Dev",
                emails_restricted_to_domains=False,
                email_address_visibility=Realm.EMAIL_ADDRESS_VISIBILITY_ADMINS,
                description=
                "The Zulip development environment default organization."
                "  It's great for testing!",
                invite_required=False,
                plan_type=Realm.SELF_HOSTED,
                org_type=Realm.ORG_TYPES["business"]["id"],
            )
            RealmDomain.objects.create(realm=zulip_realm, domain="zulip.com")
            assert zulip_realm.notifications_stream is not None
            zulip_realm.notifications_stream.name = "Verona"
            zulip_realm.notifications_stream.description = "A city in Italy"
            zulip_realm.notifications_stream.save(
                update_fields=["name", "description"])

            if options["test_suite"]:
                mit_realm = do_create_realm(
                    string_id="zephyr",
                    name="MIT",
                    emails_restricted_to_domains=True,
                    invite_required=False,
                    plan_type=Realm.SELF_HOSTED,
                    org_type=Realm.ORG_TYPES["business"]["id"],
                )
                RealmDomain.objects.create(realm=mit_realm, domain="mit.edu")

                lear_realm = do_create_realm(
                    string_id="lear",
                    name="Lear & Co.",
                    emails_restricted_to_domains=False,
                    invite_required=False,
                    plan_type=Realm.SELF_HOSTED,
                    org_type=Realm.ORG_TYPES["business"]["id"],
                )

                # Default to allowing all members to send mentions in
                # large streams for the test suite to keep
                # mention-related tests simple.
                zulip_realm.wildcard_mention_policy = Realm.WILDCARD_MENTION_POLICY_MEMBERS
                zulip_realm.save(update_fields=["wildcard_mention_policy"])

            # Create test Users (UserProfiles are automatically created,
            # as are subscriptions to the ability to receive personals).
            names = [
                ("Zoe", "*****@*****.**"),
                ("Othello, the Moor of Venice", "*****@*****.**"),
                ("Iago", "*****@*****.**"),
                ("Prospero from The Tempest", "*****@*****.**"),
                ("Cordelia, Lear's daughter", "*****@*****.**"),
                ("King Hamlet", "*****@*****.**"),
                ("aaron", "*****@*****.**"),
                ("Polonius", "*****@*****.**"),
                ("Desdemona", "*****@*****.**"),
                ("शिव", "*****@*****.**"),
            ]

            # For testing really large batches:
            # Create extra users with semi realistic names to make search
            # functions somewhat realistic.  We'll still create 1000 users
            # like Extra222 User for some predicability.
            num_names = options["extra_users"]
            num_boring_names = 300

            for i in range(min(num_names, num_boring_names)):
                full_name = f"Extra{i:03} User"
                names.append((full_name, f"extrauser{i}@zulip.com"))

            if num_names > num_boring_names:
                fnames = [
                    "Amber",
                    "Arpita",
                    "Bob",
                    "Cindy",
                    "Daniela",
                    "Dan",
                    "Dinesh",
                    "Faye",
                    "François",
                    "George",
                    "Hank",
                    "Irene",
                    "James",
                    "Janice",
                    "Jenny",
                    "Jill",
                    "John",
                    "Kate",
                    "Katelyn",
                    "Kobe",
                    "Lexi",
                    "Manish",
                    "Mark",
                    "Matt",
                    "Mayna",
                    "Michael",
                    "Pete",
                    "Peter",
                    "Phil",
                    "Phillipa",
                    "Preston",
                    "Sally",
                    "Scott",
                    "Sandra",
                    "Steve",
                    "Stephanie",
                    "Vera",
                ]
                mnames = ["de", "van", "von", "Shaw", "T."]
                lnames = [
                    "Adams",
                    "Agarwal",
                    "Beal",
                    "Benson",
                    "Bonita",
                    "Davis",
                    "George",
                    "Harden",
                    "James",
                    "Jones",
                    "Johnson",
                    "Jordan",
                    "Lee",
                    "Leonard",
                    "Singh",
                    "Smith",
                    "Patel",
                    "Towns",
                    "Wall",
                ]
                non_ascii_names = [
                    "Günter",
                    "أحمد",
                    "Magnús",
                    "आशी",
                    "イツキ",
                    "语嫣",
                    "அருண்",
                    "Александр",
                    "José",
                ]
                # to imitate emoji insertions in usernames
                raw_emojis = ["😎", "😂", "🐱‍👤"]

            for i in range(num_boring_names, num_names):
                fname = random.choice(fnames) + str(i)
                full_name = fname
                if random.random() < 0.7:
                    if random.random() < 0.3:
                        full_name += " " + random.choice(non_ascii_names)
                    else:
                        full_name += " " + random.choice(mnames)
                    if random.random() < 0.1:
                        full_name += " {} ".format(random.choice(raw_emojis))
                    else:
                        full_name += " " + random.choice(lnames)
                email = fname.lower() + "@zulip.com"
                names.append((full_name, email))

            create_users(zulip_realm, names, tos_version=settings.TOS_VERSION)

            iago = get_user_by_delivery_email("*****@*****.**", zulip_realm)
            do_change_user_role(iago,
                                UserProfile.ROLE_REALM_ADMINISTRATOR,
                                acting_user=None)
            iago.is_staff = True
            iago.save(update_fields=["is_staff"])

            desdemona = get_user_by_delivery_email("*****@*****.**",
                                                   zulip_realm)
            do_change_user_role(desdemona,
                                UserProfile.ROLE_REALM_OWNER,
                                acting_user=None)

            shiva = get_user_by_delivery_email("*****@*****.**", zulip_realm)
            do_change_user_role(shiva,
                                UserProfile.ROLE_MODERATOR,
                                acting_user=None)

            guest_user = get_user_by_delivery_email("*****@*****.**",
                                                    zulip_realm)
            guest_user.role = UserProfile.ROLE_GUEST
            guest_user.save(update_fields=["role"])

            # These bots are directly referenced from code and thus
            # are needed for the test suite.
            zulip_realm_bots = [
                ("Zulip Error Bot", "*****@*****.**"),
                ("Zulip Default Bot", "*****@*****.**"),
            ]
            for i in range(options["extra_bots"]):
                zulip_realm_bots.append(
                    (f"Extra Bot {i}", f"extrabot{i}@zulip.com"))

            create_users(zulip_realm,
                         zulip_realm_bots,
                         bot_type=UserProfile.DEFAULT_BOT)

            zoe = get_user_by_delivery_email("*****@*****.**", zulip_realm)
            zulip_webhook_bots = [
                ("Zulip Webhook Bot", "*****@*****.**"),
            ]
            # If a stream is not supplied in the webhook URL, the webhook
            # will (in some cases) send the notification as a PM to the
            # owner of the webhook bot, so bot_owner can't be None
            create_users(
                zulip_realm,
                zulip_webhook_bots,
                bot_type=UserProfile.INCOMING_WEBHOOK_BOT,
                bot_owner=zoe,
            )
            aaron = get_user_by_delivery_email("*****@*****.**", zulip_realm)

            zulip_outgoing_bots = [
                ("Outgoing Webhook", "*****@*****.**"),
            ]
            create_users(
                zulip_realm,
                zulip_outgoing_bots,
                bot_type=UserProfile.OUTGOING_WEBHOOK_BOT,
                bot_owner=aaron,
            )
            outgoing_webhook = get_user("*****@*****.**",
                                        zulip_realm)
            add_service(
                "outgoing-webhook",
                user_profile=outgoing_webhook,
                interface=Service.GENERIC,
                base_url="http://127.0.0.1:5002",
                token=generate_api_key(),
            )

            # Add the realm internal bots to each realm.
            create_if_missing_realm_internal_bots()

            # Create public streams.
            signups_stream = Realm.INITIAL_PRIVATE_STREAM_NAME

            stream_list = [
                "Verona",
                "Denmark",
                "Scotland",
                "Venice",
                "Rome",
                signups_stream,
            ]
            stream_dict: Dict[str, Dict[str, Any]] = {
                "Denmark": {
                    "description": "A Scandinavian country"
                },
                "Scotland": {
                    "description": "Located in the United Kingdom"
                },
                "Venice": {
                    "description": "A northeastern Italian city"
                },
                "Rome": {
                    "description": "Yet another Italian city",
                    "is_web_public": True
                },
            }

            bulk_create_streams(zulip_realm, stream_dict)
            recipient_streams: List[int] = [
                Stream.objects.get(name=name, realm=zulip_realm).id
                for name in stream_list
            ]

            # 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). For the test suite, we have a fixed list of
            # subscriptions to make sure test data is consistent
            # across platforms.

            subscriptions_list: List[Tuple[UserProfile, Recipient]] = []
            profiles: Sequence[UserProfile] = (
                UserProfile.objects.select_related().filter(
                    is_bot=False).order_by("email"))

            if options["test_suite"]:
                subscriptions_map = {
                    "*****@*****.**": ["Verona"],
                    "*****@*****.**": ["Verona"],
                    "*****@*****.**": ["Verona", "Denmark", signups_stream],
                    "*****@*****.**": [
                        "Verona",
                        "Denmark",
                        "Scotland",
                        signups_stream,
                    ],
                    "*****@*****.**": ["Verona", "Denmark", "Scotland"],
                    "*****@*****.**":
                    ["Verona", "Denmark", "Scotland", "Venice"],
                    "*****@*****.**":
                    ["Verona", "Denmark", "Scotland", "Venice", "Rome"],
                    "*****@*****.**": ["Verona"],
                    "*****@*****.**": [
                        "Verona",
                        "Denmark",
                        "Venice",
                        signups_stream,
                    ],
                    "*****@*****.**": ["Verona", "Denmark", "Scotland"],
                }

                for profile in profiles:
                    email = profile.delivery_email
                    if email not in subscriptions_map:
                        raise Exception(
                            f"Subscriptions not listed for user {email}")

                    for stream_name in subscriptions_map[email]:
                        stream = Stream.objects.get(name=stream_name,
                                                    realm=zulip_realm)
                        r = Recipient.objects.get(type=Recipient.STREAM,
                                                  type_id=stream.id)
                        subscriptions_list.append((profile, r))
            else:
                num_streams = len(recipient_streams)
                num_users = len(profiles)
                for i, profile in enumerate(profiles):
                    # Subscribe to some streams.
                    fraction = float(i) / num_users
                    num_recips = int(num_streams * fraction) + 1

                    for type_id in recipient_streams[:num_recips]:
                        r = Recipient.objects.get(type=Recipient.STREAM,
                                                  type_id=type_id)
                        subscriptions_list.append((profile, r))

            subscriptions_to_add: List[Subscription] = []
            event_time = timezone_now()
            all_subscription_logs: (List[RealmAuditLog]) = []

            i = 0
            for profile, recipient in subscriptions_list:
                i += 1
                color = STREAM_ASSIGNMENT_COLORS[i %
                                                 len(STREAM_ASSIGNMENT_COLORS)]
                s = Subscription(
                    recipient=recipient,
                    user_profile=profile,
                    is_user_active=profile.is_active,
                    color=color,
                )

                subscriptions_to_add.append(s)

                log = RealmAuditLog(
                    realm=profile.realm,
                    modified_user=profile,
                    modified_stream_id=recipient.type_id,
                    event_last_message_id=0,
                    event_type=RealmAuditLog.SUBSCRIPTION_CREATED,
                    event_time=event_time,
                )
                all_subscription_logs.append(log)

            Subscription.objects.bulk_create(subscriptions_to_add)
            RealmAuditLog.objects.bulk_create(all_subscription_logs)

            # Create custom profile field data
            phone_number = try_add_realm_custom_profile_field(
                zulip_realm,
                "Phone number",
                CustomProfileField.SHORT_TEXT,
                hint="")
            biography = try_add_realm_custom_profile_field(
                zulip_realm,
                "Biography",
                CustomProfileField.LONG_TEXT,
                hint="What are you known for?",
            )
            favorite_food = try_add_realm_custom_profile_field(
                zulip_realm,
                "Favorite food",
                CustomProfileField.SHORT_TEXT,
                hint="Or drink, if you'd prefer",
            )
            field_data: ProfileFieldData = {
                "vim": {
                    "text": "Vim",
                    "order": "1"
                },
                "emacs": {
                    "text": "Emacs",
                    "order": "2"
                },
            }
            favorite_editor = try_add_realm_custom_profile_field(
                zulip_realm,
                "Favorite editor",
                CustomProfileField.SELECT,
                field_data=field_data)
            birthday = try_add_realm_custom_profile_field(
                zulip_realm, "Birthday", CustomProfileField.DATE)
            favorite_website = try_add_realm_custom_profile_field(
                zulip_realm,
                "Favorite website",
                CustomProfileField.URL,
                hint="Or your personal blog's URL",
            )
            mentor = try_add_realm_custom_profile_field(
                zulip_realm, "Mentor", CustomProfileField.USER)
            github_profile = try_add_realm_default_custom_profile_field(
                zulip_realm, "github")

            # Fill in values for Iago and Hamlet
            hamlet = get_user_by_delivery_email("*****@*****.**",
                                                zulip_realm)
            do_update_user_custom_profile_data_if_changed(
                iago,
                [
                    {
                        "id": phone_number.id,
                        "value": "+1-234-567-8901"
                    },
                    {
                        "id": biography.id,
                        "value": "Betrayer of Othello."
                    },
                    {
                        "id": favorite_food.id,
                        "value": "Apples"
                    },
                    {
                        "id": favorite_editor.id,
                        "value": "emacs"
                    },
                    {
                        "id": birthday.id,
                        "value": "2000-01-01"
                    },
                    {
                        "id": favorite_website.id,
                        "value": "https://zulip.readthedocs.io/en/latest/"
                    },
                    {
                        "id": mentor.id,
                        "value": [hamlet.id]
                    },
                    {
                        "id": github_profile.id,
                        "value": "zulip"
                    },
                ],
            )
            do_update_user_custom_profile_data_if_changed(
                hamlet,
                [
                    {
                        "id": phone_number.id,
                        "value": "+0-11-23-456-7890"
                    },
                    {
                        "id":
                        biography.id,
                        "value":
                        "I am:\n* The prince of Denmark\n* Nephew to the usurping Claudius",
                    },
                    {
                        "id": favorite_food.id,
                        "value": "Dark chocolate"
                    },
                    {
                        "id": favorite_editor.id,
                        "value": "vim"
                    },
                    {
                        "id": birthday.id,
                        "value": "1900-01-01"
                    },
                    {
                        "id": favorite_website.id,
                        "value": "https://blog.zulig.org"
                    },
                    {
                        "id": mentor.id,
                        "value": [iago.id]
                    },
                    {
                        "id": github_profile.id,
                        "value": "zulipbot"
                    },
                ],
            )
        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] = list(
            UserProfile.objects.filter(is_bot=False))

        # Create a test realm emoji.
        IMAGE_FILE_PATH = static_path("images/test-images/checkbox.png")
        with open(IMAGE_FILE_PATH, "rb") as fp:
            check_add_realm_emoji(zulip_realm, "green_tick", iago, fp)

        if not options["test_suite"]:
            # Populate users with some bar data
            for user in user_profiles:
                status: int = UserPresence.ACTIVE
                date = timezone_now()
                client = get_client("website")
                if user.full_name[0] <= "H":
                    client = get_client("ZulipAndroid")
                UserPresence.objects.get_or_create(
                    user_profile=user,
                    realm_id=user.realm_id,
                    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"])
        ]

        create_alert_words(zulip_realm.id)

        # Generate a new set of test data.
        create_test_data()

        # prepopulate the URL preview/embed data for the links present
        # in the config.generate_data.json data set.  This makes it
        # possible for populate_db to run happily without Internet
        # access.
        with open("zerver/tests/fixtures/docs_url_preview_data.json",
                  "rb") as f:
            urls_with_preview_data = orjson.loads(f.read())
            for url in urls_with_preview_data:
                cache_set(url, urls_with_preview_data[url], PREVIEW_CACHE_NAME)

        if options["delete"]:
            if options["test_suite"]:
                # Create test users; the MIT ones are needed to test
                # the Zephyr mirroring codepaths.
                testsuite_mit_users = [
                    ("Fred Sipb (MIT)", "*****@*****.**"),
                    ("Athena Consulting Exchange User (MIT)",
                     "*****@*****.**"),
                    ("Esp Classroom (MIT)", "*****@*****.**"),
                ]
                create_users(mit_realm,
                             testsuite_mit_users,
                             tos_version=settings.TOS_VERSION)

                testsuite_lear_users = [
                    ("King Lear", "*****@*****.**"),
                    ("Cordelia, Lear's daughter", "*****@*****.**"),
                ]
                create_users(lear_realm,
                             testsuite_lear_users,
                             tos_version=settings.TOS_VERSION)

            if not options["test_suite"]:
                # 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

                # to imitate emoji insertions in stream names
                raw_emojis = ["😎", "😂", "🐱‍👤"]

                zulip_stream_dict: Dict[str, Dict[str, Any]] = {
                    "devel": {
                        "description": "For developing"
                    },
                    # ビデオゲーム - VideoGames (japanese)
                    "ビデオゲーム": {
                        "description":
                        "Share your favorite video games!  {}".format(
                            raw_emojis[2])
                    },
                    "announce": {
                        "description": "For announcements",
                        "stream_post_policy": Stream.STREAM_POST_POLICY_ADMINS,
                    },
                    "design": {
                        "description": "For design"
                    },
                    "support": {
                        "description": "For support"
                    },
                    "social": {
                        "description": "For socializing"
                    },
                    "test": {
                        "description": "For testing `code`"
                    },
                    "errors": {
                        "description": "For errors"
                    },
                    # 조리법 - Recipes (Korean) , Пельмени - Dumplings (Russian)
                    "조리법 " + raw_emojis[0]: {
                        "description":
                        "Everything cooking, from pasta to Пельмени"
                    },
                }

                extra_stream_names = [
                    "802.11a",
                    "Ad Hoc Network",
                    "Augmented Reality",
                    "Cycling",
                    "DPI",
                    "FAQ",
                    "FiFo",
                    "commits",
                    "Control panel",
                    "desktop",
                    "компьютеры",
                    "Data security",
                    "desktop",
                    "काम",
                    "discussions",
                    "Cloud storage",
                    "GCI",
                    "Vaporware",
                    "Recent Trends",
                    "issues",
                    "live",
                    "Health",
                    "mobile",
                    "空間",
                    "provision",
                    "hidrógeno",
                    "HR",
                    "アニメ",
                ]

                # Add stream names and stream descriptions
                for i in range(options["extra_streams"]):
                    extra_stream_name = random.choice(
                        extra_stream_names) + " " + str(i)

                    # to imitate emoji insertions in stream names
                    if random.random() <= 0.15:
                        extra_stream_name += random.choice(raw_emojis)

                    zulip_stream_dict[extra_stream_name] = {
                        "description": "Auto-generated extra stream.",
                    }

                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
                subscribe_users_to_streams(zulip_realm, zulip_stream_dict)

            create_user_groups()

            if not options["test_suite"]:
                # We populate the analytics database here for
                # development purpose only
                call_command("populate_analytics_db")

        threads = options["threads"]
        jobs: List[Tuple[int, List[List[int]], Dict[str, Any],
                         Callable[[str], int], 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,
                         random.randint(0, 10**10)))

        for job in jobs:
            generate_and_send_messages(job)

        if options["delete"]:
            if not options["test_suite"]:
                # These bots are not needed by the test suite
                # Also, we don't want interacting with each other
                # in dev setup.
                internal_zulip_users_nosubs = [
                    ("Zulip Commit Bot", "*****@*****.**"),
                    ("Zulip Trac Bot", "*****@*****.**"),
                    ("Zulip Nagios Bot", "*****@*****.**"),
                ]
                create_users(zulip_realm,
                             internal_zulip_users_nosubs,
                             bot_type=UserProfile.DEFAULT_BOT)

            mark_all_messages_as_read()
            self.stdout.write("Successfully populated test database.\n")
Example #3
0
    def handle(self, **options: 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,
                description=
                "The Zulip development environment default organization."
                "  It's great for testing!",
                invite_required=False,
                org_type=Realm.CORPORATE)
            RealmDomain.objects.create(realm=zulip_realm, domain="zulip.com")
            if options["test_suite"]:
                mit_realm = Realm.objects.create(string_id="zephyr",
                                                 name="MIT",
                                                 restricted_to_domain=True,
                                                 invite_required=False,
                                                 org_type=Realm.CORPORATE)
                RealmDomain.objects.create(realm=mit_realm, domain="mit.edu")

                lear_realm = Realm.objects.create(string_id="lear",
                                                  name="Lear & Co.",
                                                  restricted_to_domain=False,
                                                  invite_required=False,
                                                  org_type=Realm.CORPORATE)

            # Create test Users (UserProfiles are automatically created,
            # as are subscriptions to the ability to receive personals).
            names = [
                ("Zoe", "*****@*****.**"),
                ("Othello, the Moor of Venice", "*****@*****.**"),
                ("Iago", "*****@*****.**"),
                ("Prospero from The Tempest", "*****@*****.**"),
                ("Cordelia Lear", "*****@*****.**"),
                ("King Hamlet", "*****@*****.**"),
                ("aaron", "*****@*****.**"),
                ("Polonius", "*****@*****.**"),
            ]
            for i in range(options["extra_users"]):
                names.append(
                    ('Extra User %d' % (i, ), '*****@*****.**' % (i, )))
            create_users(zulip_realm, names)

            iago = get_user("*****@*****.**", zulip_realm)
            do_change_is_admin(iago, True)
            iago.is_staff = True
            iago.save(update_fields=['is_staff'])

            guest_user = get_user("*****@*****.**", zulip_realm)
            guest_user.is_guest = True
            guest_user.save(update_fields=['is_guest'])

            # 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", "*****@*****.**"),
                ("Zulip Error Bot", "*****@*****.**"),
                ("Zulip Default Bot", "*****@*****.**"),
                ("Welcome Bot", "*****@*****.**"),
            ]

            for i in range(options["extra_bots"]):
                zulip_realm_bots.append(
                    ('Extra Bot %d' % (i, ), '*****@*****.**' % (i, )))
            zulip_realm_bots.extend(all_realm_bots)
            create_users(zulip_realm,
                         zulip_realm_bots,
                         bot_type=UserProfile.DEFAULT_BOT)

            # Initialize the email gateway bot as an API Super User
            email_gateway_bot = get_system_bot(settings.EMAIL_GATEWAY_BOT)
            email_gateway_bot.is_api_super_user = True
            email_gateway_bot.save()

            zoe = get_user("*****@*****.**", zulip_realm)
            zulip_webhook_bots = [
                ("Zulip Webhook Bot", "*****@*****.**"),
            ]
            # If a stream is not supplied in the webhook URL, the webhook
            # will (in some cases) send the notification as a PM to the
            # owner of the webhook bot, so bot_owner can't be None
            create_users(zulip_realm,
                         zulip_webhook_bots,
                         bot_type=UserProfile.INCOMING_WEBHOOK_BOT,
                         bot_owner=zoe)
            aaron = get_user("*****@*****.**", zulip_realm)
            zulip_outgoing_bots = [("Outgoing Webhook",
                                    "*****@*****.**")]
            create_users(zulip_realm,
                         zulip_outgoing_bots,
                         bot_type=UserProfile.OUTGOING_WEBHOOK_BOT,
                         bot_owner=aaron)
            # TODO: Clean up this initial bot creation code
            Service.objects.create(
                name="test",
                user_profile=get_user("*****@*****.**",
                                      zulip_realm),
                base_url="http://127.0.0.1:5002/bots/followup",
                token="abcd1234",
                interface=1)

            # Create public streams.
            stream_list = ["Verona", "Denmark", "Scotland", "Venice", "Rome"]
            stream_dict = {
                "Verona": {
                    "description": "A city in Italy",
                    "invite_only": False,
                    "is_web_public": False
                },
                "Denmark": {
                    "description": "A Scandinavian country",
                    "invite_only": False,
                    "is_web_public": False
                },
                "Scotland": {
                    "description": "Located in the United Kingdom",
                    "invite_only": False,
                    "is_web_public": False
                },
                "Venice": {
                    "description": "A northeastern Italian city",
                    "invite_only": False,
                    "is_web_public": False
                },
                "Rome": {
                    "description": "Yet another Italian city",
                    "invite_only": False,
                    "is_web_public": True
                }
            }  # type: Dict[str, Dict[str, 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]
            event_time = timezone_now()
            all_subscription_logs = []  # type: (List[RealmAuditLog])
            profiles = UserProfile.objects.select_related().filter(
                is_bot=False, is_guest=False).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)

                    log = RealmAuditLog(realm=profile.realm,
                                        modified_user=profile,
                                        modified_stream_id=type_id,
                                        event_last_message_id=0,
                                        event_type='subscription_created',
                                        event_time=event_time)
                    all_subscription_logs.append(log)

            Subscription.objects.bulk_create(subscriptions_to_add)
            RealmAuditLog.objects.bulk_create(all_subscription_logs)

            # Create custom profile field data
            phone_number = try_add_realm_custom_profile_field(
                zulip_realm,
                "Phone number",
                CustomProfileField.SHORT_TEXT,
                hint='')
            biography = try_add_realm_custom_profile_field(
                zulip_realm,
                "Biography",
                CustomProfileField.LONG_TEXT,
                hint='What are you known for?')
            favorite_food = try_add_realm_custom_profile_field(
                zulip_realm,
                "Favorite food",
                CustomProfileField.SHORT_TEXT,
                hint="Or drink, if you'd prefer")
            field_data = {
                'vim': {
                    'text': 'Vim',
                    'order': '1'
                },
                'emacs': {
                    'text': 'Emacs',
                    'order': '2'
                },
            }
            favorite_editor = try_add_realm_custom_profile_field(
                zulip_realm,
                "Favorite editor",
                CustomProfileField.CHOICE,
                field_data=field_data)
            birthday = try_add_realm_custom_profile_field(
                zulip_realm, "Birthday", CustomProfileField.DATE)
            favorite_website = try_add_realm_custom_profile_field(
                zulip_realm,
                "GitHub profile",
                CustomProfileField.URL,
                hint="Or your personal blog's URL")

            # Fill in values for Iago and Hamlet
            hamlet = get_user("*****@*****.**", zulip_realm)
            do_update_user_custom_profile_data(iago, [
                {
                    "id": phone_number.id,
                    "value": "+1-234-567-8901"
                },
                {
                    "id": biography.id,
                    "value": "Betrayer of Othello."
                },
                {
                    "id": favorite_food.id,
                    "value": "Apples"
                },
                {
                    "id": favorite_editor.id,
                    "value": "emacs"
                },
                {
                    "id": birthday.id,
                    "value": "2000-1-1"
                },
                {
                    "id": favorite_website.id,
                    "value": "https://github.com/zulip/zulip"
                },
            ])
            do_update_user_custom_profile_data(hamlet, [
                {
                    "id": phone_number.id,
                    "value": "+0-11-23-456-7890"
                },
                {
                    "id": biography.id,
                    "value": "Prince of Denmark, and other things!"
                },
                {
                    "id": favorite_food.id,
                    "value": "Dark chocolate"
                },
                {
                    "id": favorite_editor.id,
                    "value": "vim"
                },
                {
                    "id": birthday.id,
                    "value": "1900-1-1"
                },
                {
                    "id": favorite_website.id,
                    "value": "https://blog.zulig.org"
                },
            ])
        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.filter(
            is_bot=False))  # type: List[UserProfile]

        # Create a test realm emoji.
        IMAGE_FILE_PATH = os.path.join(settings.STATIC_ROOT, 'images',
                                       'test-images', 'checkbox.png')
        with open(IMAGE_FILE_PATH, 'rb') as fp:
            check_add_realm_emoji(zulip_realm, 'green_tick', iago, fp)

        if not options["test_suite"]:
            # Populate users with some bar data
            for user in user_profiles:
                status = UserPresence.ACTIVE  # type: int
                date = timezone_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"])
        ]

        # Generate a new set of test data.
        create_test_data()

        # prepopulate the URL preview/embed data for the links present
        # in the config.generate_data.json data set.  This makes it
        # possible for populate_db to run happily without Internet
        # access.
        with open("zerver/tests/fixtures/docs_url_preview_data.json",
                  "r") as f:
            urls_with_preview_data = ujson.load(f)
            for url in urls_with_preview_data:
                cache_set(url, urls_with_preview_data[url], PREVIEW_CACHE_NAME)

        threads = options["threads"]
        jobs = [
        ]  # type: List[Tuple[int, List[List[int]], Dict[str, Any], Callable[[str], int], 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,
                         random.randint(0, 10**10)))

        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)", "*****@*****.**"),
                    ("Athena Consulting Exchange User (MIT)",
                     "*****@*****.**"),
                    ("Esp Classroom (MIT)", "*****@*****.**"),
                ]
                create_users(mit_realm, testsuite_mit_users)

                testsuite_lear_users = [
                    ("King Lear", "*****@*****.**"),
                    ("Cordelia Lear", "*****@*****.**"),
                ]
                create_users(lear_realm, testsuite_lear_users)

            if not options["test_suite"]:
                # 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,
                        "is_web_public": False
                    },
                    "all": {
                        "description": "For everything",
                        "invite_only": False,
                        "is_web_public": False
                    },
                    "announce": {
                        "description": "For announcements",
                        "invite_only": False,
                        "is_web_public": False
                    },
                    "design": {
                        "description": "For design",
                        "invite_only": False,
                        "is_web_public": False
                    },
                    "support": {
                        "description": "For support",
                        "invite_only": False,
                        "is_web_public": False
                    },
                    "social": {
                        "description": "For socializing",
                        "invite_only": False,
                        "is_web_public": False
                    },
                    "test": {
                        "description": "For testing",
                        "invite_only": False,
                        "is_web_public": False
                    },
                    "errors": {
                        "description": "For errors",
                        "invite_only": False,
                        "is_web_public": False
                    },
                    "sales": {
                        "description": "For sales discussion",
                        "invite_only": False,
                        "is_web_public": False
                    }
                }  # type: Dict[str, Dict[str, Any]]

                # Calculate the maximum number of digits in any extra stream's
                # number, since a stream with name "Extra Stream 3" could show
                # up after "Extra Stream 29". (Used later to pad numbers with
                # 0s).
                maximum_digits = len(str(options['extra_streams'] - 1))

                for i in range(options['extra_streams']):
                    # Pad the number with 0s based on `maximum_digits`.
                    number_str = str(i).zfill(maximum_digits)

                    extra_stream_name = 'Extra Stream ' + number_str

                    zulip_stream_dict[extra_stream_name] = {
                        "description": "Auto-generated extra stream.",
                        "invite_only": False,
                        "is_web_public": False,
                    }

                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 = []
                event_time = timezone_now()
                all_subscription_logs = []
                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)

                        log = RealmAuditLog(realm=profile.realm,
                                            modified_user=profile,
                                            modified_stream=stream,
                                            event_last_message_id=0,
                                            event_type='subscription_created',
                                            event_time=event_time)
                        all_subscription_logs.append(log)
                Subscription.objects.bulk_create(subscriptions_to_add)
                RealmAuditLog.objects.bulk_create(all_subscription_logs)

                # These bots are not needed by the test suite
                internal_zulip_users_nosubs = [
                    ("Zulip Commit Bot", "*****@*****.**"),
                    ("Zulip Trac Bot", "*****@*****.**"),
                    ("Zulip Nagios Bot", "*****@*****.**"),
                ]
                create_users(zulip_realm,
                             internal_zulip_users_nosubs,
                             bot_type=UserProfile.DEFAULT_BOT)

            zulip_cross_realm_bots = [
                ("Zulip Feedback Bot", "*****@*****.**"),
            ]
            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)

            if not options["test_suite"]:
                # Update pointer of each user to point to the last message in their
                # UserMessage rows with sender_id=user_profile_id.
                users = list(
                    UserMessage.objects.filter(message__sender_id=F(
                        'user_profile_id')).values('user_profile_id').annotate(
                            pointer=Max('message_id')))
                for user in users:
                    UserProfile.objects.filter(
                        id=user['user_profile_id']).update(
                            pointer=user['pointer'])

            create_user_groups()
            self.stdout.write("Successfully populated test database.\n")
Example #4
0
    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,
                description=
                "The Zulip development environment default organization."
                "  It's great for testing!",
                invite_required=False,
                org_type=Realm.CORPORATE)
            RealmDomain.objects.create(realm=zulip_realm, domain="zulip.com")
            if options["test_suite"]:
                mit_realm = Realm.objects.create(string_id="zephyr",
                                                 name="MIT",
                                                 restricted_to_domain=True,
                                                 invite_required=False,
                                                 org_type=Realm.CORPORATE)
                RealmDomain.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", "*****@*****.**"),
                ("Othello, the Moor of Venice", "*****@*****.**"),
                ("Iago", "*****@*****.**"),
                ("Prospero from The Tempest", "*****@*****.**"),
                ("Cordelia Lear", "*****@*****.**"),
                ("King Hamlet", "*****@*****.**"),
                ("aaron", "*****@*****.**"),
            ]
            for i in range(options["extra_users"]):
                names.append(
                    ('Extra User %d' % (i, ), '*****@*****.**' % (i, )))
            create_users(zulip_realm, names)

            iago = get_user("*****@*****.**", zulip_realm)
            do_change_is_admin(iago, True)
            iago.is_staff = True
            iago.save(update_fields=['is_staff'])

            # 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", "*****@*****.**"),
                ("Zulip Error Bot", "*****@*****.**"),
                ("Zulip Default Bot", "*****@*****.**"),
                ("Welcome Bot", "*****@*****.**"),
            ]

            for i in range(options["extra_bots"]):
                zulip_realm_bots.append(
                    ('Extra Bot %d' % (i, ), '*****@*****.**' % (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", "*****@*****.**"),
            ]
            create_users(zulip_realm,
                         zulip_webhook_bots,
                         bot_type=UserProfile.INCOMING_WEBHOOK_BOT)
            zulip_outgoing_bots = [("Outgoing Webhook",
                                    "*****@*****.**")]
            aaron = get_user("*****@*****.**", zulip_realm)
            create_users(zulip_realm,
                         zulip_outgoing_bots,
                         bot_type=UserProfile.OUTGOING_WEBHOOK_BOT,
                         bot_owner=aaron)
            # TODO: Clean up this initial bot creation code
            Service.objects.create(
                name="test",
                user_profile=get_user("*****@*****.**",
                                      zulip_realm),
                base_url="http://127.0.0.1:5002/bots/followup",
                token="abcd1234",
                interface=1)

            # 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]
            event_time = timezone_now()
            all_subscription_logs = []  # type: (List[RealmAuditLog])
            profiles = UserProfile.objects.select_related().filter(
                is_bot=False).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)

                    log = RealmAuditLog(realm=profile.realm,
                                        modified_user=profile,
                                        modified_stream_id=type_id,
                                        event_last_message_id=0,
                                        event_type='subscription_created',
                                        event_time=event_time)
                    all_subscription_logs.append(log)

            Subscription.objects.bulk_create(subscriptions_to_add)
            RealmAuditLog.objects.bulk_create(all_subscription_logs)
        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.filter(
            is_bot=False))  # type: List[UserProfile]

        # Create a test realm emoji.
        IMAGE_FILE_PATH = os.path.join(settings.STATIC_ROOT, 'images',
                                       'test-images', 'checkbox.png')
        UPLOADED_EMOJI_FILE_NAME = 'green_tick.png'
        with open(IMAGE_FILE_PATH, 'rb') as fp:
            upload_backend.upload_emoji_image(fp, UPLOADED_EMOJI_FILE_NAME,
                                              iago)
            RealmEmoji.objects.create(
                name='green_tick',
                author=iago,
                realm=zulip_realm,
                deactivated=False,
                file_name=UPLOADED_EMOJI_FILE_NAME,
            )

        if not options["test_suite"]:
            # Populate users with some bar data
            for user in user_profiles:
                status = UserPresence.ACTIVE  # type: int
                date = timezone_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"])
        ]

        # Generate a new set of test data.
        create_test_data()

        threads = options["threads"]
        jobs = [
        ]  # type: List[Tuple[int, List[List[int]], Dict[str, Any], Callable[[str], int], 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,
                         random.randint(0, 10**10)))

        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)", "*****@*****.**"),
                    ("Athena Consulting Exchange User (MIT)",
                     "*****@*****.**"),
                    ("Esp Classroom (MIT)", "*****@*****.**"),
                ]
                create_users(mit_realm, testsuite_mit_users)

            if not options["test_suite"]:
                # Initialize the email gateway bot as an API Super User
                email_gateway_bot = get_system_bot(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 = []
                event_time = timezone_now()
                all_subscription_logs = []
                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)

                        log = RealmAuditLog(realm=profile.realm,
                                            modified_user=profile,
                                            modified_stream=stream,
                                            event_last_message_id=0,
                                            event_type='subscription_created',
                                            event_time=event_time)
                        all_subscription_logs.append(log)
                Subscription.objects.bulk_create(subscriptions_to_add)
                RealmAuditLog.objects.bulk_create(all_subscription_logs)

                # These bots are not needed by the test suite
                internal_zulip_users_nosubs = [
                    ("Zulip Commit Bot", "*****@*****.**"),
                    ("Zulip Trac Bot", "*****@*****.**"),
                    ("Zulip Nagios Bot", "*****@*****.**"),
                ]
                create_users(zulip_realm,
                             internal_zulip_users_nosubs,
                             bot_type=UserProfile.DEFAULT_BOT)

            zulip_cross_realm_bots = [
                ("Zulip Feedback Bot", "*****@*****.**"),
            ]
            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)

            if not options["test_suite"]:
                # Update pointer of each user to point to the last message in their
                # UserMessage rows with sender_id=user_profile_id.
                users = list(
                    UserMessage.objects.filter(message__sender_id=F(
                        'user_profile_id')).values('user_profile_id').annotate(
                            pointer=Max('message_id')))
                for user in users:
                    UserProfile.objects.filter(
                        id=user['user_profile_id']).update(
                            pointer=user['pointer'])

            create_user_groups()
            self.stdout.write("Successfully populated test database.\n")
Example #5
0
    def handle(self, **options: Any) -> None:
        if options["percent_huddles"] + options["percent_personals"] > 100:
            self.stderr.write("Error!  More than 100% of messages allocated.\n")
            return

        # Get consistent data for backend tests.
        if options["test_suite"]:
            random.seed(0)

        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", emails_restricted_to_domains=True,
                description="The Zulip development environment default organization."
                            "  It's great for testing!",
                invite_required=False, org_type=Realm.CORPORATE)
            RealmDomain.objects.create(realm=zulip_realm, domain="zulip.com")
            if options["test_suite"]:
                mit_realm = Realm.objects.create(
                    string_id="zephyr", name="MIT", emails_restricted_to_domains=True,
                    invite_required=False, org_type=Realm.CORPORATE)
                RealmDomain.objects.create(realm=mit_realm, domain="mit.edu")

                lear_realm = Realm.objects.create(
                    string_id="lear", name="Lear & Co.", emails_restricted_to_domains=False,
                    invite_required=False, org_type=Realm.CORPORATE)

            # Create test Users (UserProfiles are automatically created,
            # as are subscriptions to the ability to receive personals).
            names = [
                ("Zoe", "*****@*****.**"),
                ("Othello, the Moor of Venice", "*****@*****.**"),
                ("Iago", "*****@*****.**"),
                ("Prospero from The Tempest", "*****@*****.**"),
                ("Cordelia Lear", "*****@*****.**"),
                ("King Hamlet", "*****@*****.**"),
                ("aaron", "*****@*****.**"),
                ("Polonius", "*****@*****.**"),
            ]
            for i in range(options["extra_users"]):
                names.append(('Extra User %d' % (i,), '*****@*****.**' % (i,)))
            create_users(zulip_realm, names)

            iago = get_user("*****@*****.**", zulip_realm)
            do_change_is_admin(iago, True)
            iago.is_staff = True
            iago.save(update_fields=['is_staff'])

            guest_user = get_user("*****@*****.**", zulip_realm)
            guest_user.is_guest = True
            guest_user.save(update_fields=['is_guest'])

            # 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", "*****@*****.**"),
                ("Zulip Error Bot", "*****@*****.**"),
                ("Zulip Default Bot", "*****@*****.**"),
                ("Welcome Bot", "*****@*****.**"),
            ]

            for i in range(options["extra_bots"]):
                zulip_realm_bots.append(('Extra Bot %d' % (i,), '*****@*****.**' % (i,)))
            zulip_realm_bots.extend(all_realm_bots)
            create_users(zulip_realm, zulip_realm_bots, bot_type=UserProfile.DEFAULT_BOT)

            # Initialize the email gateway bot as an API Super User
            email_gateway_bot = get_system_bot(settings.EMAIL_GATEWAY_BOT)
            email_gateway_bot.is_api_super_user = True
            email_gateway_bot.save()

            zoe = get_user("*****@*****.**", zulip_realm)
            zulip_webhook_bots = [
                ("Zulip Webhook Bot", "*****@*****.**"),
            ]
            # If a stream is not supplied in the webhook URL, the webhook
            # will (in some cases) send the notification as a PM to the
            # owner of the webhook bot, so bot_owner can't be None
            create_users(zulip_realm, zulip_webhook_bots,
                         bot_type=UserProfile.INCOMING_WEBHOOK_BOT, bot_owner=zoe)
            aaron = get_user("*****@*****.**", zulip_realm)

            zulip_outgoing_bots = [
                ("Outgoing Webhook", "*****@*****.**")
            ]
            create_users(zulip_realm, zulip_outgoing_bots,
                         bot_type=UserProfile.OUTGOING_WEBHOOK_BOT, bot_owner=aaron)
            outgoing_webhook = get_user("*****@*****.**", zulip_realm)
            add_service("outgoing-webhook", user_profile=outgoing_webhook, interface=Service.GENERIC,
                        base_url="http://127.0.0.1:5002", token=generate_api_key())

            # Add the realm internl bots to each realm.
            create_if_missing_realm_internal_bots()

            # Create public streams.
            stream_list = ["Verona", "Denmark", "Scotland", "Venice", "Rome"]
            stream_dict = {
                "Verona": {"description": "A city in Italy"},
                "Denmark": {"description": "A Scandinavian country"},
                "Scotland": {"description": "Located in the United Kingdom"},
                "Venice": {"description": "A northeastern Italian city"},
                "Rome": {"description": "Yet another Italian city", "is_web_public": True}
            }  # type: Dict[str, Dict[str, 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). For the test suite, we have a fixed list of
            # subscriptions to make sure test data is consistent
            # across platforms.

            subscriptions_list = []  # type: List[Tuple[UserProfile, Recipient]]
            profiles = UserProfile.objects.select_related().filter(
                is_bot=False).order_by("email")  # type: Sequence[UserProfile]

            if options["test_suite"]:
                subscriptions_map = {
                    '*****@*****.**': ['Verona'],
                    '*****@*****.**': ['Verona'],
                    '*****@*****.**': ['Verona', 'Denmark'],
                    '*****@*****.**': ['Verona', 'Denmark', 'Scotland'],
                    '*****@*****.**': ['Verona', 'Denmark', 'Scotland'],
                    '*****@*****.**': ['Verona', 'Denmark', 'Scotland', 'Venice'],
                    '*****@*****.**': ['Verona', 'Denmark', 'Scotland', 'Venice', 'Rome'],
                    '*****@*****.**': ['Verona'],
                }

                for profile in profiles:
                    if profile.email not in subscriptions_map:
                        raise Exception('Subscriptions not listed for user %s' % (profile.email,))

                    for stream_name in subscriptions_map[profile.email]:
                        stream = Stream.objects.get(name=stream_name)
                        r = Recipient.objects.get(type=Recipient.STREAM, type_id=stream.id)
                        subscriptions_list.append((profile, r))
            else:
                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)
                        subscriptions_list.append((profile, r))

            subscriptions_to_add = []  # type: List[Subscription]
            event_time = timezone_now()
            all_subscription_logs = []  # type: (List[RealmAuditLog])

            i = 0
            for profile, recipient in subscriptions_list:
                i += 1
                color = STREAM_ASSIGNMENT_COLORS[i % len(STREAM_ASSIGNMENT_COLORS)]
                s = Subscription(
                    recipient=recipient,
                    user_profile=profile,
                    color=color)

                subscriptions_to_add.append(s)

                log = RealmAuditLog(realm=profile.realm,
                                    modified_user=profile,
                                    modified_stream_id=recipient.type_id,
                                    event_last_message_id=0,
                                    event_type=RealmAuditLog.SUBSCRIPTION_CREATED,
                                    event_time=event_time)
                all_subscription_logs.append(log)

            Subscription.objects.bulk_create(subscriptions_to_add)
            RealmAuditLog.objects.bulk_create(all_subscription_logs)

            # Create custom profile field data
            phone_number = try_add_realm_custom_profile_field(zulip_realm, "Phone number",
                                                              CustomProfileField.SHORT_TEXT,
                                                              hint='')
            biography = try_add_realm_custom_profile_field(zulip_realm, "Biography",
                                                           CustomProfileField.LONG_TEXT,
                                                           hint='What are you known for?')
            favorite_food = try_add_realm_custom_profile_field(zulip_realm, "Favorite food",
                                                               CustomProfileField.SHORT_TEXT,
                                                               hint="Or drink, if you'd prefer")
            field_data = {
                'vim': {'text': 'Vim', 'order': '1'},
                'emacs': {'text': 'Emacs', 'order': '2'},
            }
            favorite_editor = try_add_realm_custom_profile_field(zulip_realm,
                                                                 "Favorite editor",
                                                                 CustomProfileField.CHOICE,
                                                                 field_data=field_data)
            birthday = try_add_realm_custom_profile_field(zulip_realm, "Birthday",
                                                          CustomProfileField.DATE)
            favorite_website = try_add_realm_custom_profile_field(zulip_realm, "GitHub profile",
                                                                  CustomProfileField.URL,
                                                                  hint="Or your personal blog's URL")
            mentor = try_add_realm_custom_profile_field(zulip_realm, "Mentor",
                                                        CustomProfileField.USER)

            # Fill in values for Iago and Hamlet
            hamlet = get_user("*****@*****.**", zulip_realm)
            do_update_user_custom_profile_data(iago, [
                {"id": phone_number.id, "value": "+1-234-567-8901"},
                {"id": biography.id, "value": "Betrayer of Othello."},
                {"id": favorite_food.id, "value": "Apples"},
                {"id": favorite_editor.id, "value": "emacs"},
                {"id": birthday.id, "value": "2000-1-1"},
                {"id": favorite_website.id, "value": "https://github.com/zulip/zulip"},
                {"id": mentor.id, "value": [hamlet.id]},
            ])
            do_update_user_custom_profile_data(hamlet, [
                {"id": phone_number.id, "value": "+0-11-23-456-7890"},
                {"id": biography.id, "value": "Prince of Denmark, and other things!"},
                {"id": favorite_food.id, "value": "Dark chocolate"},
                {"id": favorite_editor.id, "value": "vim"},
                {"id": birthday.id, "value": "1900-1-1"},
                {"id": favorite_website.id, "value": "https://blog.zulig.org"},
                {"id": mentor.id, "value": [iago.id]},
            ])
        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.filter(is_bot=False))  # type: List[UserProfile]

        # Create a test realm emoji.
        IMAGE_FILE_PATH = os.path.join(settings.STATIC_ROOT, 'images', 'test-images', 'checkbox.png')
        with open(IMAGE_FILE_PATH, 'rb') as fp:
            check_add_realm_emoji(zulip_realm, 'green_tick', iago, fp)

        if not options["test_suite"]:
            # Populate users with some bar data
            for user in user_profiles:
                status = UserPresence.ACTIVE  # type: int
                date = timezone_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"])]

        # Generate a new set of test data.
        create_test_data()

        # prepopulate the URL preview/embed data for the links present
        # in the config.generate_data.json data set.  This makes it
        # possible for populate_db to run happily without Internet
        # access.
        with open("zerver/tests/fixtures/docs_url_preview_data.json", "r") as f:
            urls_with_preview_data = ujson.load(f)
            for url in urls_with_preview_data:
                cache_set(url, urls_with_preview_data[url], PREVIEW_CACHE_NAME)

        threads = options["threads"]
        jobs = []  # type: List[Tuple[int, List[List[int]], Dict[str, Any], Callable[[str], int], 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, random.randint(0, 10**10)))

        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)", "*****@*****.**"),
                    ("Athena Consulting Exchange User (MIT)", "*****@*****.**"),
                    ("Esp Classroom (MIT)", "*****@*****.**"),
                ]
                create_users(mit_realm, testsuite_mit_users)

                testsuite_lear_users = [
                    ("King Lear", "*****@*****.**"),
                    ("Cordelia Lear", "*****@*****.**"),
                ]
                create_users(lear_realm, testsuite_lear_users)

            if not options["test_suite"]:
                # 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"},
                    "all": {"description": "For everything"},
                    "announce": {"description": "For announcements", 'is_announcement_only': True},
                    "design": {"description": "For design"},
                    "support": {"description": "For support"},
                    "social": {"description": "For socializing"},
                    "test": {"description": "For testing"},
                    "errors": {"description": "For errors"},
                    "sales": {"description": "For sales discussion"}
                }  # type: Dict[str, Dict[str, Any]]

                # Calculate the maximum number of digits in any extra stream's
                # number, since a stream with name "Extra Stream 3" could show
                # up after "Extra Stream 29". (Used later to pad numbers with
                # 0s).
                maximum_digits = len(str(options['extra_streams'] - 1))

                for i in range(options['extra_streams']):
                    # Pad the number with 0s based on `maximum_digits`.
                    number_str = str(i).zfill(maximum_digits)

                    extra_stream_name = 'Extra Stream ' + number_str

                    zulip_stream_dict[extra_stream_name] = {
                        "description": "Auto-generated extra stream.",
                    }

                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
                subscribe_users_to_streams(zulip_realm, zulip_stream_dict)

                # These bots are not needed by the test suite
                internal_zulip_users_nosubs = [
                    ("Zulip Commit Bot", "*****@*****.**"),
                    ("Zulip Trac Bot", "*****@*****.**"),
                    ("Zulip Nagios Bot", "*****@*****.**"),
                ]
                create_users(zulip_realm, internal_zulip_users_nosubs, bot_type=UserProfile.DEFAULT_BOT)

            zulip_cross_realm_bots = [
                ("Zulip Feedback Bot", "*****@*****.**"),
            ]
            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)

            if not options["test_suite"]:
                # Update pointer of each user to point to the last message in their
                # UserMessage rows with sender_id=user_profile_id.
                users = list(UserMessage.objects.filter(
                    message__sender_id=F('user_profile_id')).values(
                    'user_profile_id').annotate(pointer=Max('message_id')))
                for user in users:
                    UserProfile.objects.filter(id=user['user_profile_id']).update(
                        pointer=user['pointer'])

            create_user_groups()

            if not options["test_suite"]:
                # We populate the analytics database here for
                # development purpose only
                call_command('populate_analytics_db')
            self.stdout.write("Successfully populated test database.\n")
Example #6
0
    def handle(self, **options: 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,
                description="The Zulip development environment default organization."
                            "  It's great for testing!",
                invite_required=False, org_type=Realm.CORPORATE)
            RealmDomain.objects.create(realm=zulip_realm, domain="zulip.com")
            if options["test_suite"]:
                mit_realm = Realm.objects.create(
                    string_id="zephyr", name="MIT", restricted_to_domain=True,
                    invite_required=False, org_type=Realm.CORPORATE)
                RealmDomain.objects.create(realm=mit_realm, domain="mit.edu")

                lear_realm = Realm.objects.create(
                    string_id="lear", name="Lear & Co.", restricted_to_domain=False,
                    invite_required=False, org_type=Realm.CORPORATE)

            # Create test Users (UserProfiles are automatically created,
            # as are subscriptions to the ability to receive personals).
            names = [
                ("Zoe", "*****@*****.**"),
                ("Othello, the Moor of Venice", "*****@*****.**"),
                ("Iago", "*****@*****.**"),
                ("Prospero from The Tempest", "*****@*****.**"),
                ("Cordelia Lear", "*****@*****.**"),
                ("King Hamlet", "*****@*****.**"),
                ("aaron", "*****@*****.**"),
            ]
            for i in range(options["extra_users"]):
                names.append(('Extra User %d' % (i,), '*****@*****.**' % (i,)))
            create_users(zulip_realm, names)

            iago = get_user("*****@*****.**", zulip_realm)
            do_change_is_admin(iago, True)
            iago.is_staff = True
            iago.save(update_fields=['is_staff'])

            # 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", "*****@*****.**"),
                ("Zulip Error Bot", "*****@*****.**"),
                ("Zulip Default Bot", "*****@*****.**"),
                ("Welcome Bot", "*****@*****.**"),
            ]

            for i in range(options["extra_bots"]):
                zulip_realm_bots.append(('Extra Bot %d' % (i,), '*****@*****.**' % (i,)))
            zulip_realm_bots.extend(all_realm_bots)
            create_users(zulip_realm, zulip_realm_bots, bot_type=UserProfile.DEFAULT_BOT)

            zoe = get_user("*****@*****.**", zulip_realm)
            zulip_webhook_bots = [
                ("Zulip Webhook Bot", "*****@*****.**"),
            ]
            # If a stream is not supplied in the webhook URL, the webhook
            # will (in some cases) send the notification as a PM to the
            # owner of the webhook bot, so bot_owner can't be None
            create_users(zulip_realm, zulip_webhook_bots,
                         bot_type=UserProfile.INCOMING_WEBHOOK_BOT, bot_owner=zoe)
            aaron = get_user("*****@*****.**", zulip_realm)
            zulip_outgoing_bots = [
                ("Outgoing Webhook", "*****@*****.**")
            ]
            create_users(zulip_realm, zulip_outgoing_bots,
                         bot_type=UserProfile.OUTGOING_WEBHOOK_BOT, bot_owner=aaron)
            # TODO: Clean up this initial bot creation code
            Service.objects.create(
                name="test",
                user_profile=get_user("*****@*****.**", zulip_realm),
                base_url="http://127.0.0.1:5002/bots/followup",
                token="abcd1234",
                interface=1)

            # 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]
            event_time = timezone_now()
            all_subscription_logs = []  # type: (List[RealmAuditLog])
            profiles = UserProfile.objects.select_related().filter(
                is_bot=False).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)

                    log = RealmAuditLog(realm=profile.realm,
                                        modified_user=profile,
                                        modified_stream_id=type_id,
                                        event_last_message_id=0,
                                        event_type='subscription_created',
                                        event_time=event_time)
                    all_subscription_logs.append(log)

            Subscription.objects.bulk_create(subscriptions_to_add)
            RealmAuditLog.objects.bulk_create(all_subscription_logs)

            if not options['test_suite']:
                # Create custom profile field data
                phone_number = try_add_realm_custom_profile_field(zulip_realm, "Phone number",
                                                                  CustomProfileField.SHORT_TEXT)
                biography = try_add_realm_custom_profile_field(zulip_realm, "Biography",
                                                               CustomProfileField.LONG_TEXT)
                favorite_integer = try_add_realm_custom_profile_field(zulip_realm, "Favorite integer",
                                                                      CustomProfileField.INTEGER)

                # Fill in values for Iago and Hamlet
                hamlet = get_user("*****@*****.**", zulip_realm)
                do_update_user_custom_profile_data(iago, [
                    {"id": phone_number.id, "value": "+1-234-567-8901"},
                    {"id": biography.id, "value": "Betrayer of Othello."},
                    {"id": favorite_integer.id, "value": 17},
                ])
                do_update_user_custom_profile_data(hamlet, [
                    {"id": phone_number.id, "value": "+0-11-23-456-7890"},
                    {"id": biography.id, "value": "Prince of Denmark, and other things!"},
                    {"id": favorite_integer.id, "value": 12},
                ])
        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.filter(is_bot=False))  # type: List[UserProfile]

        # Create a test realm emoji.
        IMAGE_FILE_PATH = os.path.join(settings.STATIC_ROOT, 'images', 'test-images', 'checkbox.png')
        UPLOADED_EMOJI_FILE_NAME = 'green_tick.png'
        with open(IMAGE_FILE_PATH, 'rb') as fp:
            upload_backend.upload_emoji_image(fp, UPLOADED_EMOJI_FILE_NAME, iago)
            RealmEmoji.objects.create(
                name='green_tick',
                author=iago,
                realm=zulip_realm,
                deactivated=False,
                file_name=UPLOADED_EMOJI_FILE_NAME,
            )

        if not options["test_suite"]:
            # Populate users with some bar data
            for user in user_profiles:
                status = UserPresence.ACTIVE  # type: int
                date = timezone_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"])]

        # Generate a new set of test data.
        create_test_data()

        threads = options["threads"]
        jobs = []  # type: List[Tuple[int, List[List[int]], Dict[str, Any], Callable[[str], int], 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, random.randint(0, 10**10)))

        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)", "*****@*****.**"),
                    ("Athena Consulting Exchange User (MIT)", "*****@*****.**"),
                    ("Esp Classroom (MIT)", "*****@*****.**"),
                ]
                create_users(mit_realm, testsuite_mit_users)

                testsuite_lear_users = [
                    ("King Lear", "*****@*****.**"),
                    ("Cordelia Lear", "*****@*****.**"),
                ]
                create_users(lear_realm, testsuite_lear_users)

            if not options["test_suite"]:
                # Initialize the email gateway bot as an API Super User
                email_gateway_bot = get_system_bot(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]]

                # Calculate the maximum number of digits in any extra stream's
                # number, since a stream with name "Extra Stream 3" could show
                # up after "Extra Stream 29". (Used later to pad numbers with
                # 0s).
                maximum_digits = len(str(options['extra_streams'] - 1))

                for i in range(options['extra_streams']):
                    # Pad the number with 0s based on `maximum_digits`.
                    number_str = str(i).zfill(maximum_digits)

                    extra_stream_name = 'Extra Stream ' + number_str

                    zulip_stream_dict[extra_stream_name] = {
                        "description": "Auto-generated extra stream.",
                        "invite_only": False,
                    }

                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 = []
                event_time = timezone_now()
                all_subscription_logs = []
                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)

                        log = RealmAuditLog(realm=profile.realm,
                                            modified_user=profile,
                                            modified_stream=stream,
                                            event_last_message_id=0,
                                            event_type='subscription_created',
                                            event_time=event_time)
                        all_subscription_logs.append(log)
                Subscription.objects.bulk_create(subscriptions_to_add)
                RealmAuditLog.objects.bulk_create(all_subscription_logs)

                # These bots are not needed by the test suite
                internal_zulip_users_nosubs = [
                    ("Zulip Commit Bot", "*****@*****.**"),
                    ("Zulip Trac Bot", "*****@*****.**"),
                    ("Zulip Nagios Bot", "*****@*****.**"),
                ]
                create_users(zulip_realm, internal_zulip_users_nosubs, bot_type=UserProfile.DEFAULT_BOT)

            zulip_cross_realm_bots = [
                ("Zulip Feedback Bot", "*****@*****.**"),
            ]
            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)

            if not options["test_suite"]:
                # Update pointer of each user to point to the last message in their
                # UserMessage rows with sender_id=user_profile_id.
                users = list(UserMessage.objects.filter(
                    message__sender_id=F('user_profile_id')).values(
                    'user_profile_id').annotate(pointer=Max('message_id')))
                for user in users:
                    UserProfile.objects.filter(id=user['user_profile_id']).update(
                        pointer=user['pointer'])

            create_user_groups()
            self.stdout.write("Successfully populated test database.\n")
Example #7
0
def create_streams(realms, realm, stream_list):
    # type: (Mapping[text_type, Realm], Realm, Iterable[text_type]) -> None
    stream_set = set() # type: Set[Tuple[text_type, text_type]]
    for stream_name in stream_list:
        stream_set.add((realm.domain, stream_name))
    bulk_create_streams(realms, stream_set)
Example #8
0
def restore_saved_messages():
    # type: () -> None
    old_messages = [] # type: List[Dict[str, Any]]
    duplicate_suppression_hash = {} # type: Dict[str, bool]

    stream_dict = {} # type: Dict[Tuple[text_type, text_type], Tuple[text_type, text_type]]
    user_set = set() # type: Set[Tuple[text_type, text_type, text_type, bool]]
    email_set = set([u.email for u in UserProfile.objects.all()]) # type: Set[text_type]
    realm_set = set() # type: Set[text_type]
    # Initial client_set is nonempty temporarily because we don't have
    # clients in logs at all right now -- later we can start with nothing.
    client_set = set(["populate_db", "website", "zephyr_mirror"])
    huddle_user_set = set() # type: Set[Tuple[text_type, ...]]
    # First, determine all the objects our messages will need.
    print(datetime.datetime.now(), "Creating realms/streams/etc...")
    def process_line(line):
        # type: (str) -> None
        old_message_json = line.strip()

        # Due to populate_db's shakespeare mode, we have a lot of
        # duplicate messages in our log that only differ in their
        # logged ID numbers (same timestamp, content, etc.).  With
        # sqlite, bulk creating those messages won't work properly: in
        # particular, the first 100 messages will actually only result
        # in 20 rows ending up in the target table, which screws up
        # the below accounting where for handling changing
        # subscriptions, we assume that the Nth row populate_db
        # created goes with the Nth non-subscription row of the input
        # So suppress the duplicates when using sqlite.
        if "sqlite" in settings.DATABASES["default"]["ENGINE"]:
            tmp_message = ujson.loads(old_message_json)
            tmp_message['id'] = '1'
            duplicate_suppression_key = ujson.dumps(tmp_message)
            if duplicate_suppression_key in duplicate_suppression_hash:
                return
            duplicate_suppression_hash[duplicate_suppression_key] = True

        old_message = ujson.loads(old_message_json)
        message_type = old_message["type"]

        # Lower case emails and domains; it will screw up
        # deduplication if we don't
        def fix_email(email):
            # type: (text_type) -> text_type
            return email.strip().lower()

        if message_type in ["stream", "huddle", "personal"]:
            old_message["sender_email"] = fix_email(old_message["sender_email"])
            # Fix the length on too-long messages before we start processing them
            if len(old_message["content"]) > MAX_MESSAGE_LENGTH:
                old_message["content"] = "[ This message was deleted because it was too long ]"
        if message_type in ["subscription_added", "subscription_removed"]:
            old_message["domain"] = old_message["domain"].lower()
            old_message["user"] = fix_email(old_message["user"])
        elif message_type == "subscription_property":
            old_message["user"] = fix_email(old_message["user"])
        elif message_type == "user_email_changed":
            old_message["old_email"] = fix_email(old_message["old_email"])
            old_message["new_email"] = fix_email(old_message["new_email"])
        elif message_type.startswith("user_"):
            old_message["user"] = fix_email(old_message["user"])
        elif message_type.startswith("enable_"):
            old_message["user"] = fix_email(old_message["user"])

        if message_type == 'personal':
            old_message["recipient"][0]["email"] = fix_email(old_message["recipient"][0]["email"])
        elif message_type == "huddle":
            for i in range(len(old_message["recipient"])):
                old_message["recipient"][i]["email"] = fix_email(old_message["recipient"][i]["email"])

        old_messages.append(old_message)

        if message_type in ["subscription_added", "subscription_removed"]:
            stream_name = old_message["name"].strip() # type: text_type
            canon_stream_name = stream_name.lower()
            if canon_stream_name not in stream_dict:
                stream_dict[(old_message["domain"], canon_stream_name)] = \
                    (old_message["domain"], stream_name)
        elif message_type == "user_created":
            user_set.add((old_message["user"], old_message["full_name"], old_message["short_name"], False))
        elif message_type == "realm_created":
            realm_set.add(old_message["domain"])

        if message_type not in ["stream", "huddle", "personal"]:
            return

        sender_email = old_message["sender_email"]

        domain = split_email_to_domain(sender_email)
        realm_set.add(domain)

        if old_message["sender_email"] not in email_set:
            user_set.add((old_message["sender_email"],
                          old_message["sender_full_name"],
                          old_message["sender_short_name"],
                          False))

        if 'sending_client' in old_message:
            client_set.add(old_message['sending_client'])

        if message_type == 'stream':
            stream_name = old_message["recipient"].strip()
            canon_stream_name = stream_name.lower()
            if canon_stream_name not in stream_dict:
                stream_dict[(domain, canon_stream_name)] = (domain, stream_name)
        elif message_type == 'personal':
            u = old_message["recipient"][0]
            if u["email"] not in email_set:
                user_set.add((u["email"], u["full_name"], u["short_name"], False))
                email_set.add(u["email"])
        elif message_type == 'huddle':
            for u in old_message["recipient"]:
                user_set.add((u["email"], u["full_name"], u["short_name"], False))
                if u["email"] not in email_set:
                    user_set.add((u["email"], u["full_name"], u["short_name"], False))
                    email_set.add(u["email"])
            huddle_user_set.add(tuple(sorted(set(u["email"] for u in old_message["recipient"]))))
        else:
            raise ValueError('Bad message type')

    event_glob = os.path.join(settings.EVENT_LOG_DIR, 'events.*')
    for filename in sorted(glob.glob(event_glob)):
        with open(filename, "r") as message_log:
            for line in message_log.readlines():
                process_line(line)

    stream_recipients = {} # type: Dict[Tuple[int, text_type], Recipient]
    user_recipients = {} # type: Dict[text_type, Recipient]
    huddle_recipients = {} # type: Dict[text_type, Recipient]

    # Then, create the objects our messages need.
    print(datetime.datetime.now(), "Creating realms...")
    bulk_create_realms(realm_set)

    realms = {} # type: Dict[text_type, Realm]
    for realm in Realm.objects.all():
        realms[realm.domain] = realm

    print(datetime.datetime.now(), "Creating clients...")
    bulk_create_clients(client_set)

    clients = {} # type: Dict[text_type, Client]
    for client in Client.objects.all():
        clients[client.name] = client

    print(datetime.datetime.now(), "Creating streams...")
    bulk_create_streams(realms, list(stream_dict.values()))

    streams = {} # type: Dict[int, Stream]
    for stream in Stream.objects.all():
        streams[stream.id] = stream
    for recipient in Recipient.objects.filter(type=Recipient.STREAM):
        stream_recipients[(streams[recipient.type_id].realm_id,
                           streams[recipient.type_id].name.lower())] = recipient

    print(datetime.datetime.now(), "Creating users...")
    bulk_create_users(realms, user_set, tos_version=settings.TOS_VERSION)

    users = {} # type: Dict[text_type, UserProfile]
    users_by_id = {} # type: Dict[int, UserProfile]
    for user_profile in UserProfile.objects.select_related().all():
        users[user_profile.email] = user_profile
        users_by_id[user_profile.id] = user_profile
    for recipient in Recipient.objects.filter(type=Recipient.PERSONAL):
        user_recipients[users_by_id[recipient.type_id].email] = recipient

    print(datetime.datetime.now(), "Creating huddles...")
    bulk_create_huddles(users, huddle_user_set)

    huddles_by_id = {} # type: Dict[int, Huddle]
    for huddle in Huddle.objects.all():
        huddles_by_id[huddle.id] = huddle
    for recipient in Recipient.objects.filter(type=Recipient.HUDDLE):
        huddle_recipients[huddles_by_id[recipient.type_id].huddle_hash] = recipient

    # TODO: Add a special entry type in the log that is a subscription
    # change and import those as we go to make subscription changes
    # take effect!
    print(datetime.datetime.now(), "Importing subscriptions...")
    subscribers = {} # type: Dict[int, Set[int]]
    for s in Subscription.objects.select_related().all():
        if s.active:
            subscribers.setdefault(s.recipient.id, set()).add(s.user_profile.id)

    # Then create all the messages, without talking to the DB!
    print(datetime.datetime.now(), "Importing messages, part 1...")
    first_message_id = None
    if Message.objects.exists():
        first_message_id = Message.objects.all().order_by("-id")[0].id + 1

    messages_to_create = [] # type: List[Message]
    for idx, old_message in enumerate(old_messages):
        message_type = old_message["type"]
        if message_type not in ["stream", "huddle", "personal"]:
            continue

        message = Message()

        sender_email = old_message["sender_email"]
        domain = split_email_to_domain(sender_email)
        realm = realms[domain]

        message.sender = users[sender_email]
        type_hash = {"stream": Recipient.STREAM,
                     "huddle": Recipient.HUDDLE,
                     "personal": Recipient.PERSONAL}

        if 'sending_client' in old_message:
            message.sending_client = clients[old_message['sending_client']]
        elif sender_email in ["*****@*****.**", "*****@*****.**", "*****@*****.**",
                              "*****@*****.**", "*****@*****.**"]:
            message.sending_client = clients['populate_db']
        elif realm.domain == "zulip.com":
            message.sending_client = clients["website"]
        elif realm.domain == "mit.edu":
            message.sending_client = clients['zephyr_mirror']
        else:
            message.sending_client = clients['populate_db']

        message.type = type_hash[message_type]
        message.content = old_message["content"]
        message.subject = old_message["subject"]
        message.pub_date = timestamp_to_datetime(old_message["timestamp"])

        if message.type == Recipient.PERSONAL:
            message.recipient = user_recipients[old_message["recipient"][0]["email"]]
        elif message.type == Recipient.STREAM:
            message.recipient = stream_recipients[(realm.id,
                                                   old_message["recipient"].lower())]
        elif message.type == Recipient.HUDDLE:
            huddle_hash = get_huddle_hash([users[u["email"]].id
                                           for u in old_message["recipient"]])
            message.recipient = huddle_recipients[huddle_hash]
        else:
            raise ValueError('Bad message type')
        messages_to_create.append(message)

    print(datetime.datetime.now(), "Importing messages, part 2...")
    Message.objects.bulk_create(messages_to_create)
    messages_to_create = []

    # Finally, create all the UserMessage objects
    print(datetime.datetime.now(), "Importing usermessages, part 1...")
    personal_recipients = {} # type: Dict[int, bool]
    for r in Recipient.objects.filter(type = Recipient.PERSONAL):
        personal_recipients[r.id] = True

    all_messages = Message.objects.all() # type: Sequence[Message]
    user_messages_to_create = [] # type: List[UserMessage]

    messages_by_id = {} # type: Dict[int, Message]
    for message in all_messages:
        messages_by_id[message.id] = message

    if len(messages_by_id) == 0:
        print(datetime.datetime.now(), "No old messages to replay")
        return

    if first_message_id is None:
        first_message_id = min(messages_by_id.keys())

    tot_user_messages = 0
    pending_subs = {} # type: Dict[Tuple[int, int], bool]
    current_message_id = first_message_id
    pending_colors = {} # type: Dict[Tuple[text_type, text_type], text_type]
    for old_message in old_messages:
        message_type = old_message["type"]
        if message_type == 'subscription_added':
            stream_key = (realms[old_message["domain"]].id, old_message["name"].strip().lower())
            subscribers.setdefault(stream_recipients[stream_key].id,
                                   set()).add(users[old_message["user"]].id)
            pending_subs[(stream_recipients[stream_key].id,
                          users[old_message["user"]].id)] = True
            continue
        elif message_type == "subscription_removed":
            stream_key = (realms[old_message["domain"]].id, old_message["name"].strip().lower())
            user_id = users[old_message["user"]].id
            subscribers.setdefault(stream_recipients[stream_key].id, set())
            try:
                subscribers[stream_recipients[stream_key].id].remove(user_id)
            except KeyError:
                print("Error unsubscribing %s from %s: not subscribed" % (
                    old_message["user"], old_message["name"]))
            pending_subs[(stream_recipients[stream_key].id,
                          users[old_message["user"]].id)] = False
            continue
        elif message_type == "user_activated" or message_type == "user_created":
            # These are rare, so just handle them the slow way
            user_profile = users[old_message["user"]]
            join_date = timestamp_to_datetime(old_message['timestamp'])
            do_activate_user(user_profile, log=False, join_date=join_date)
            # Update the cache of users to show this user as activated
            users_by_id[user_profile.id] = user_profile
            users[old_message["user"]] = user_profile
            continue
        elif message_type == "user_deactivated":
            user_profile = users[old_message["user"]]
            do_deactivate_user(user_profile, log=False)
            continue
        elif message_type == "user_change_password":
            # Just handle these the slow way
            user_profile = users[old_message["user"]]
            do_change_password(user_profile, old_message["pwhash"], log=False,
                               hashed_password=True)
            continue
        elif message_type == "user_change_full_name":
            # Just handle these the slow way
            user_profile = users[old_message["user"]]
            user_profile.full_name = old_message["full_name"]
            user_profile.save(update_fields=["full_name"])
            continue
        elif message_type == "enable_desktop_notifications_changed":
            # Just handle these the slow way
            user_profile = users[old_message["user"]]
            user_profile.enable_desktop_notifications = (old_message["enable_desktop_notifications"] != "false")
            user_profile.save(update_fields=["enable_desktop_notifications"])
            continue
        elif message_type == "enable_sounds_changed":
            user_profile = users[old_message["user"]]
            user_profile.enable_sounds = (old_message["enable_sounds"] != "false")
            user_profile.save(update_fields=["enable_sounds"])
        elif message_type == "enable_offline_email_notifications_changed":
            user_profile = users[old_message["user"]]
            user_profile.enable_offline_email_notifications = (
                old_message["enable_offline_email_notifications"] != "false")
            user_profile.save(update_fields=["enable_offline_email_notifications"])
            continue
        elif message_type == "enable_offline_push_notifications_changed":
            user_profile = users[old_message["user"]]
            user_profile.enable_offline_push_notifications = (
                old_message["enable_offline_push_notifications"] != "false")
            user_profile.save(update_fields=["enable_offline_push_notifications"])
            continue
        elif message_type == "default_streams":
            set_default_streams(get_realm(old_message["domain"]),
                                old_message["streams"])
            continue
        elif message_type == "subscription_property":
            property_name = old_message.get("property")
            if property_name == "stream_color" or property_name == "color":
                color = old_message.get("color", old_message.get("value"))
                pending_colors[(old_message["user"],
                                old_message["stream_name"].lower())] = color
            elif property_name in ["in_home_view", "notifications"]:
                # TODO: Handle this
                continue
            else:
                raise RuntimeError("Unknown property %s" % (property_name,))
            continue
        elif message_type == "realm_created":
            # No action required
            continue
        elif message_type in ["user_email_changed", "update_onboarding", "update_message"]:
            # TODO: Handle these
            continue
        if message_type not in ["stream", "huddle", "personal"]:
            raise RuntimeError("Unexpected message type %s" % (message_type,))

        message = messages_by_id[current_message_id]
        current_message_id += 1

        if message.recipient_id not in subscribers:
            # Nobody received this message -- probably due to our
            # subscriptions being out-of-date.
            continue

        recipient_user_ids = set() # type: Set[int]
        for user_profile_id in subscribers[message.recipient_id]:
            recipient_user_ids.add(user_profile_id)
        if message.recipient_id in personal_recipients:
            # Include the sender in huddle recipients
            recipient_user_ids.add(message.sender_id)

        for user_profile_id in recipient_user_ids:
            if users_by_id[user_profile_id].is_active:
                um = UserMessage(user_profile_id=user_profile_id,
                                 message=message)
                user_messages_to_create.append(um)

        if len(user_messages_to_create) > 100000:
            tot_user_messages += len(user_messages_to_create)
            UserMessage.objects.bulk_create(user_messages_to_create)
            user_messages_to_create = []

    print(datetime.datetime.now(), "Importing usermessages, part 2...")
    tot_user_messages += len(user_messages_to_create)
    UserMessage.objects.bulk_create(user_messages_to_create)

    print(datetime.datetime.now(), "Finalizing subscriptions...")
    current_subs = {} # type: Dict[Tuple[int, int], bool]
    current_subs_obj = {} # type: Dict[Tuple[int, int], Subscription]
    for s in Subscription.objects.select_related().all():
        current_subs[(s.recipient_id, s.user_profile_id)] = s.active
        current_subs_obj[(s.recipient_id, s.user_profile_id)] = s

    subscriptions_to_add = [] # type: List[Subscription]
    subscriptions_to_change = [] # type: List[Tuple[Tuple[int, int], bool]]
    for pending_sub in pending_subs.keys():
        (recipient_id, user_profile_id) = pending_sub
        current_state = current_subs.get(pending_sub)
        if pending_subs[pending_sub] == current_state:
            # Already correct in the database
            continue
        elif current_state is not None:
            subscriptions_to_change.append((pending_sub, pending_subs[pending_sub]))
            continue

        s = Subscription(recipient_id=recipient_id,
                         user_profile_id=user_profile_id,
                         active=pending_subs[pending_sub])
        subscriptions_to_add.append(s)
    Subscription.objects.bulk_create(subscriptions_to_add)
    for (sub_tuple, active) in subscriptions_to_change:
        current_subs_obj[sub_tuple].active = active
        current_subs_obj[sub_tuple].save(update_fields=["active"])

    subs = {} # type: Dict[Tuple[int, int], Subscription]
    for sub in Subscription.objects.all():
        subs[(sub.user_profile_id, sub.recipient_id)] = sub

    # TODO: do restore of subscription colors -- we're currently not
    # logging changes so there's little point in having the code :(

    print(datetime.datetime.now(), "Finished importing %s messages (%s usermessages)" % \
        (len(all_messages), tot_user_messages))

    site = Site.objects.get_current()
    site.domain = 'zulip.com'
    site.save()

    print(datetime.datetime.now(), "Filling in user pointers...")

    # Set restored pointers to the very latest messages
    for user_profile in UserProfile.objects.all():
        try:
            top = UserMessage.objects.filter(
                user_profile_id=user_profile.id).order_by("-message")[0]
            user_profile.pointer = top.message_id
        except IndexError:
            user_profile.pointer = -1
        user_profile.save(update_fields=["pointer"])

    print(datetime.datetime.now(), "Done replaying old messages")