-
-
Notifications
You must be signed in to change notification settings - Fork 8.1k
/
Copy pathpopulate_db.py
1424 lines (1263 loc) · 57.6 KB
/
populate_db.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
995
996
997
998
999
1000
import itertools
import os
import random
from collections import defaultdict
from collections.abc import Mapping, Sequence
from datetime import datetime, timedelta
from typing import Any
import bmemcached
import orjson
from django.conf import settings
from django.contrib.sessions.models import Session
from django.core.files.base import File
from django.core.management import call_command
from django.core.management.base import CommandParser
from django.core.validators import validate_email
from django.db import connection
from django.db.models import F
from django.db.models.signals import post_delete
from django.utils.timezone import now as timezone_now
from typing_extensions import override
from scripts.lib.zulip_tools import get_or_create_dev_uuid_var_path
from zerver.actions.create_realm import do_create_realm
from zerver.actions.custom_profile_fields import (
do_update_user_custom_profile_data_if_changed,
try_add_realm_custom_profile_field,
try_add_realm_default_custom_profile_field,
)
from zerver.actions.message_send import build_message_send_dict, do_send_messages
from zerver.actions.realm_emoji import check_add_realm_emoji
from zerver.actions.realm_linkifiers import do_add_linkifier
from zerver.actions.scheduled_messages import check_schedule_message
from zerver.actions.streams import bulk_add_subscriptions
from zerver.actions.user_groups import create_user_group_in_database
from zerver.actions.user_settings import do_change_user_setting
from zerver.actions.users import do_change_user_role
from zerver.lib.bulk_create import bulk_create_streams
from zerver.lib.generate_test_data import create_test_data, generate_topics
from zerver.lib.management import ZulipBaseCommand
from zerver.lib.onboarding import create_if_missing_realm_internal_bots
from zerver.lib.push_notifications import logger as push_notifications_logger
from zerver.lib.remote_server import get_realms_info_for_push_bouncer
from zerver.lib.server_initialization import create_internal_realm, create_users
from zerver.lib.storage import static_path
from zerver.lib.stream_color import STREAM_ASSIGNMENT_COLORS
from zerver.lib.types import AnalyticsDataUploadLevel, ProfileFieldData
from zerver.lib.users import add_service
from zerver.lib.utils import generate_api_key
from zerver.models import (
AlertWord,
Client,
CustomProfileField,
DefaultStream,
DirectMessageGroup,
Draft,
Message,
Reaction,
Realm,
RealmAuditLog,
RealmDomain,
RealmUserDefault,
Recipient,
Service,
Stream,
Subscription,
UserMessage,
UserPresence,
UserProfile,
)
from zerver.models.alert_words import flush_alert_word
from zerver.models.clients import get_client
from zerver.models.onboarding_steps import OnboardingStep
from zerver.models.realm_audit_logs import AuditLogEventType
from zerver.models.realms import WildcardMentionPolicyEnum, get_realm
from zerver.models.recipients import get_or_create_direct_message_group
from zerver.models.streams import get_stream
from zerver.models.users import get_user, get_user_by_delivery_email, get_user_profile_by_id
from zilencer.models import RemoteRealm, RemoteZulipServer, RemoteZulipServerAuditLog
from zilencer.views import update_remote_realm_data_for_server
# Disable the push notifications bouncer to avoid enqueuing updates in
# maybe_enqueue_audit_log_upload during early setup.
settings.ZULIP_SERVICE_PUSH_NOTIFICATIONS = False
settings.ZULIP_SERVICE_SUBMIT_USAGE_STATISTICS = False
settings.ZULIP_SERVICE_SECURITY_ALERTS = False
settings.ANALYTICS_DATA_UPLOAD_LEVEL = AnalyticsDataUploadLevel.NONE
settings.USING_TORNADO = False
# Disable using memcached caches to avoid 'unsupported pickle
# protocol' errors if `populate_db` is run with a different Python
# from `run-dev`.
default_cache = settings.CACHES["default"]
settings.CACHES["default"] = {
"BACKEND": "django.core.cache.backends.locmem.LocMemCache",
}
DEFAULT_EMOJIS = [
("+1", "1f44d"),
("smiley", "1f603"),
("eyes", "1f440"),
("crying_cat_face", "1f63f"),
("arrow_up", "2b06"),
("confetti_ball", "1f38a"),
("hundred_points", "1f4af"),
]
def clear_database() -> None:
# Hacky function only for use inside populate_db. Designed to
# allow running populate_db repeatedly in series to work without
# flushing memcached or clearing the database manually.
# With `zproject.test_settings`, we aren't using real memcached
# and; we only need to flush memcached if we're populating a
# database that would be used with it (i.e. zproject.dev_settings).
if default_cache["BACKEND"] == "zerver.lib.singleton_bmemcached.SingletonBMemcached":
memcached_client = bmemcached.Client(
(default_cache["LOCATION"],),
**default_cache["OPTIONS"],
)
try:
memcached_client.flush_all()
finally:
memcached_client.disconnect_all()
model: Any = None # Hack because mypy doesn't know these are model classes
# The after-delete signal on this just updates caches, and slows
# down the deletion noticeably. Remove the signal and replace it
# after we're done.
post_delete.disconnect(flush_alert_word, sender=AlertWord)
for model in [
Message,
Stream,
AlertWord,
UserProfile,
Recipient,
Realm,
Subscription,
DirectMessageGroup,
UserMessage,
Client,
DefaultStream,
RemoteRealm,
RemoteZulipServer,
]:
model.objects.all().delete()
Session.objects.all().delete()
post_delete.connect(flush_alert_word, sender=AlertWord)
def subscribe_users_to_streams(realm: Realm, stream_dict: dict[str, dict[str, Any]]) -> None:
subscriptions_to_add = []
event_time = timezone_now()
all_subscription_logs = []
profiles = UserProfile.objects.select_related("realm").filter(realm=realm)
for i, stream_name in enumerate(stream_dict):
stream = Stream.objects.get(name=stream_name, realm=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,
is_user_active=profile.is_active,
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=AuditLogEventType.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)
def create_alert_words(realm_id: int) -> None:
user_ids = UserProfile.objects.filter(
realm_id=realm_id,
is_bot=False,
is_active=True,
).values_list("id", flat=True)
alert_words = [
"algorithms",
"complexity",
"founded",
"galaxy",
"grammar",
"illustrious",
"natural",
"objective",
"people",
"robotics",
"study",
]
recs: list[AlertWord] = []
for user_id in user_ids:
random.shuffle(alert_words)
recs.extend(
AlertWord(realm_id=realm_id, user_profile_id=user_id, word=word)
for word in alert_words[:4]
)
AlertWord.objects.bulk_create(recs)
class Command(ZulipBaseCommand):
help = "Populate a test database"
@override
def add_arguments(self, parser: CommandParser) -> None:
parser.add_argument(
"-n", "--num-messages", type=int, default=1000, help="The number of messages to create."
)
parser.add_argument(
"-o",
"--oldest-message-days",
type=int,
default=5,
help="The start of the time range where messages could have been sent.",
)
parser.add_argument(
"-b",
"--batch-size",
type=int,
default=1000,
help="How many messages to process in a single batch",
)
parser.add_argument(
"--extra-users", type=int, default=0, help="The number of extra users to create"
)
parser.add_argument(
"--extra-bots", type=int, default=0, help="The number of extra bots to create"
)
parser.add_argument(
"--extra-streams", type=int, default=0, help="The number of extra streams to create"
)
parser.add_argument("--max-topics", type=int, help="The number of maximum topics to create")
parser.add_argument(
"--direct-message-groups",
dest="num_direct_message_groups",
type=int,
default=3,
help="The number of direct message groups to create.",
)
parser.add_argument(
"--personals",
dest="num_personals",
type=int,
default=6,
help="The number of personal pairs to create.",
)
parser.add_argument("--threads", type=int, default=1, help="The number of threads to use.")
parser.add_argument(
"--percent-direct-message-groups",
type=float,
default=15,
help="The percent of messages to be direct message groups.",
)
parser.add_argument(
"--percent-personals",
type=float,
default=15,
help="The percent of messages to be personals.",
)
parser.add_argument(
"--stickiness",
type=float,
default=20,
help="The percent of messages to repeat recent folks.",
)
parser.add_argument(
"--nodelete",
action="store_false",
dest="delete",
help="Whether to delete all the existing messages.",
)
parser.add_argument(
"--test-suite",
action="store_true",
help="Configures populate_db to create a deterministic "
"data set for the backend tests.",
)
@override
def handle(self, *args: Any, **options: Any) -> None:
# Suppress spammy output from the push notifications logger
push_notifications_logger.disabled = True
if options["percent_direct_message_groups"] + 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,
# eliminating the possibility of such coincidences.
cursor.execute("SELECT setval('zerver_recipient_id_seq', 100)")
if options["max_topics"] is None:
# If max_topics is not set, we use a default that's big
# enough "show all topics" should appear, and scales slowly
# with the number of messages.
options["max_topics"] = 8 + options["num_messages"] // 1000
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.actions.create_realm.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,
description="The Zulip development environment default organization."
" It's great for testing!",
invite_required=False,
plan_type=Realm.PLAN_TYPE_SELF_HOSTED,
org_type=Realm.ORG_TYPES["business"]["id"],
enable_read_receipts=True,
enable_spectator_access=True,
)
RealmDomain.objects.create(realm=zulip_realm, domain="zulip.com")
assert zulip_realm.new_stream_announcements_stream is not None
zulip_realm.new_stream_announcements_stream.name = "Verona"
zulip_realm.new_stream_announcements_stream.description = "A city in Italy"
zulip_realm.new_stream_announcements_stream.save(update_fields=["name", "description"])
realm_user_default = RealmUserDefault.objects.get(realm=zulip_realm)
realm_user_default.enter_sends = True
realm_user_default.email_address_visibility = (
RealmUserDefault.EMAIL_ADDRESS_VISIBILITY_ADMINS
)
realm_user_default.save()
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.PLAN_TYPE_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.PLAN_TYPE_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 = WildcardMentionPolicyEnum.MEMBERS
zulip_realm.save(update_fields=["wildcard_mention_policy"])
# Realms should have matching RemoteRealm entries - simulating having realms registered
# with the bouncer, which is going to be the primary case for modern servers. Tests
# wanting to have missing registrations, or simulating legacy server scenarios,
# should delete RemoteRealms to explicit set things up.
assert isinstance(settings.ZULIP_ORG_ID, str)
assert isinstance(settings.ZULIP_ORG_KEY, str)
server = RemoteZulipServer.objects.create(
uuid=settings.ZULIP_ORG_ID,
api_key=settings.ZULIP_ORG_KEY,
hostname=settings.EXTERNAL_HOST,
last_updated=timezone_now(),
contact_email="[email protected]",
)
RemoteZulipServerAuditLog.objects.create(
event_type=AuditLogEventType.REMOTE_SERVER_CREATED,
server=server,
event_time=server.last_updated,
)
update_remote_realm_data_for_server(server, get_realms_info_for_push_bouncer())
# Create test Users (UserProfiles are automatically created,
# as are subscriptions to the ability to receive personals).
names = [
("Zoe", "[email protected]"),
("Othello, the Moor of Venice", "[email protected]"),
("Iago", "[email protected]"),
("Prospero from The Tempest", "[email protected]"),
("Cordelia, Lear's daughter", "[email protected]"),
("King Hamlet", "[email protected]"),
("aaron", "[email protected]"),
("Polonius", "[email protected]"),
("Desdemona", "[email protected]"),
("शिव", "[email protected]"),
]
# 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 predictability.
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 += f" {random.choice(raw_emojis)} "
else:
full_name += " " + random.choice(lnames)
email = fname.lower().encode("ascii", "ignore").decode("ascii") + "@zulip.com"
validate_email(email)
names.append((full_name, email))
create_users(zulip_realm, names, tos_version=settings.TERMS_OF_SERVICE_VERSION)
# Add time zones to some users. Ideally, this would be
# done in the initial create_users calls, but the
# tuple-based interface for that function doesn't support
# doing so.
def assign_time_zone_by_delivery_email(delivery_email: str, new_time_zone: str) -> None:
u = get_user_by_delivery_email(delivery_email, zulip_realm)
u.timezone = new_time_zone
u.save(update_fields=["timezone"])
# Note: Hamlet keeps default time zone of "".
assign_time_zone_by_delivery_email("[email protected]", "US/Pacific")
assign_time_zone_by_delivery_email("[email protected]", "US/Pacific")
assign_time_zone_by_delivery_email("[email protected]", "US/Eastern")
assign_time_zone_by_delivery_email("[email protected]", "US/Eastern")
assign_time_zone_by_delivery_email("[email protected]", "Canada/Newfoundland")
assign_time_zone_by_delivery_email("[email protected]", "Asia/Shanghai") # China
assign_time_zone_by_delivery_email("[email protected]", "Asia/Kolkata") # India
assign_time_zone_by_delivery_email("[email protected]", "UTC")
iago = get_user_by_delivery_email("[email protected]", zulip_realm)
do_change_user_role(iago, UserProfile.ROLE_REALM_ADMINISTRATOR, acting_user=None)
iago.is_staff = True
iago.save(update_fields=["is_staff"])
# We need to create at least two test draft for Iago for the sake
# of the cURL tests. Two since one will be deleted.
Draft.objects.create(
user_profile=iago,
recipient=None,
topic="Release Notes",
content="Release 4.0 will contain ...",
last_edit_time=timezone_now(),
)
Draft.objects.create(
user_profile=iago,
recipient=None,
topic="Release Notes",
content="Release 4.0 will contain many new features such as ... ",
last_edit_time=timezone_now(),
)
desdemona = get_user_by_delivery_email("[email protected]", zulip_realm)
do_change_user_role(desdemona, UserProfile.ROLE_REALM_OWNER, acting_user=None)
shiva = get_user_by_delivery_email("[email protected]", zulip_realm)
do_change_user_role(shiva, UserProfile.ROLE_MODERATOR, acting_user=None)
polonius = get_user_by_delivery_email("[email protected]", zulip_realm)
do_change_user_role(polonius, UserProfile.ROLE_GUEST, acting_user=None)
# These bots are directly referenced from code and thus
# are needed for the test suite.
zulip_realm_bots = [
("Zulip Default Bot", "[email protected]"),
*(
(f"Extra Bot {i}", f"extrabot{i}@zulip.com")
for i in range(options["extra_bots"])
),
]
create_users(
zulip_realm, zulip_realm_bots, bot_type=UserProfile.DEFAULT_BOT, bot_owner=desdemona
)
zoe = get_user_by_delivery_email("[email protected]", zulip_realm)
zulip_webhook_bots = [
("Zulip Webhook Bot", "[email protected]"),
]
# 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("[email protected]", zulip_realm)
zulip_outgoing_bots = [
("Outgoing Webhook", "[email protected]"),
]
create_users(
zulip_realm,
zulip_outgoing_bots,
bot_type=UserProfile.OUTGOING_WEBHOOK_BOT,
bot_owner=aaron,
)
outgoing_webhook = get_user("[email protected]", 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 streams.
zulip_discussion_channel_name = str(Realm.ZULIP_DISCUSSION_CHANNEL_NAME)
zulip_sandbox_channel_name = str(Realm.ZULIP_SANDBOX_CHANNEL_NAME)
stream_list = [
"Verona",
"Denmark",
"Scotland",
"Venice",
"Rome",
"core team",
zulip_discussion_channel_name,
zulip_sandbox_channel_name,
]
stream_dict: dict[str, dict[str, Any]] = {
"Denmark": {"description": "A Scandinavian country"},
"Scotland": {"description": "Located in the United Kingdom", "creator": iago},
"Venice": {"description": "A northeastern Italian city", "creator": polonius},
"Rome": {"description": "Yet another Italian city", "is_web_public": True},
"core team": {
"description": "A private channel for core team members",
"invite_only": True,
"creator": desdemona,
},
}
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] = list(
UserProfile.objects.select_related("realm").filter(is_bot=False).order_by("email")
)
if options["test_suite"]:
subscriptions_map = {
"[email protected]": ["Verona"],
"[email protected]": ["Verona"],
"[email protected]": [
"Verona",
"Denmark",
"core team",
zulip_discussion_channel_name,
zulip_sandbox_channel_name,
],
"[email protected]": [
"Verona",
"Denmark",
"Scotland",
"core team",
zulip_discussion_channel_name,
zulip_sandbox_channel_name,
],
"[email protected]": ["Verona", "Denmark", "Scotland"],
"[email protected]": ["Verona", "Denmark", "Scotland", "Venice"],
"[email protected]": ["Verona", "Denmark", "Scotland", "Venice", "Rome"],
"[email protected]": ["Verona"],
"[email protected]": [
"Verona",
"Denmark",
"Venice",
"core team",
zulip_discussion_channel_name,
zulip_sandbox_channel_name,
],
"[email protected]": ["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=AuditLogEventType.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 = {
"0": {"text": "Vim", "order": "1"},
"1": {"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")
pronouns = try_add_realm_custom_profile_field(
zulip_realm,
"Pronouns",
CustomProfileField.PRONOUNS,
hint="What pronouns should people use to refer to you?",
)
# Fill in values for Iago and Hamlet
hamlet = get_user_by_delivery_email("[email protected]", 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": "1"},
{"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"},
{"id": pronouns.id, "value": "he/him"},
],
)
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": "0"},
{"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"},
{"id": pronouns.id, "value": "he/him"},
],
)
# We need to create at least one scheduled message for Iago for the api-test
# cURL example to delete an existing scheduled message.
check_schedule_message(
sender=iago,
client=get_client("populate_db"),
recipient_type_name="stream",
message_to=[Stream.objects.get(name="Denmark", realm=zulip_realm).id],
topic_name="test-api",
message_content="It's time to celebrate the anniversary of provisioning this development environment :tada:!",
deliver_at=timezone_now() + timedelta(days=365),
realm=zulip_realm,
)
check_schedule_message(
sender=iago,
client=get_client("populate_db"),
recipient_type_name="private",
message_to=[iago.id],
topic_name=None,
message_content="Note to self: It's been a while since you've provisioned this development environment.",
deliver_at=timezone_now() + timedelta(days=365),
realm=zulip_realm,
)
do_add_linkifier(
zulip_realm,
"#D(?P<id>[0-9]{2,8})",
"https://github.com/zulip/zulip-desktop/pull/{id}",
acting_user=None,
)
do_add_linkifier(
zulip_realm,
"zulip-mobile#(?P<id>[0-9]{2,8})",
"https://github.com/zulip/zulip-mobile/pull/{id}",
acting_user=None,
)
do_add_linkifier(
zulip_realm,
"zulip-(?P<repo>[a-zA-Z-_0-9]+)#(?P<id>[0-9]{2,8})",
"https://github.com/zulip/{repo}/pull/{id}",
acting_user=None,
)
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, realm=zulip_realm)
)
if options["test_suite"]:
# As we plan to change the default values for 'automatically_follow_topics_policy' and
# 'automatically_unmute_topics_in_muted_streams_policy' in the future, it will lead to
# skewing a lot of our tests, which now need to take into account extra events and database queries.
#
# We explicitly set the values for both settings to 'AUTOMATICALLY_CHANGE_VISIBILITY_POLICY_NEVER'
# to make the tests independent of the default values.
#
# We have separate tests to verify events generated, database query counts,
# and other important details related to the above-mentioned settings.
#
# We set the value of 'automatically_follow_topics_where_mentioned' to 'False' so that it
# does not increase the number of events and db queries while running tests.
for user in user_profiles:
do_change_user_setting(
user,
"automatically_follow_topics_policy",
UserProfile.AUTOMATICALLY_CHANGE_VISIBILITY_POLICY_NEVER,
acting_user=None,
)
do_change_user_setting(
user,
"automatically_unmute_topics_in_muted_streams_policy",
UserProfile.AUTOMATICALLY_CHANGE_VISIBILITY_POLICY_NEVER,
acting_user=None,
)
do_change_user_setting(
user,
"automatically_follow_topics_where_mentioned",
False,
acting_user=None,
)
# 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, File(fp, name="checkbox.png"), "image/png"
)
if not options["test_suite"]:
# Populate users with some bar data
for user in user_profiles:
date = timezone_now()
UserPresence.objects.get_or_create(
user_profile=user,
realm_id=user.realm_id,
defaults={"last_active_time": date, "last_connected_time": date},
)
user_profiles_ids = []
onboarding_steps = []
for user_profile in user_profiles:
user_profiles_ids.append(user_profile.id)
onboarding_steps.append(
OnboardingStep(
user=user_profile, onboarding_step="narrow_to_dm_with_welcome_bot_new_user"
)
)
# Existing users shouldn't narrow to DM with welcome bot on first login.
OnboardingStep.objects.bulk_create(onboarding_steps)
# Create several initial direct message groups
for i in range(options["num_direct_message_groups"]):
get_or_create_direct_message_group(
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()
if options["delete"]:
if options["test_suite"]:
# Create test users; the MIT ones are needed to test
# the Zephyr mirroring codepaths.
event_time = timezone_now()
testsuite_mit_users = [
("Fred Sipb (MIT)", "[email protected]"),
("Athena Consulting Exchange User (MIT)", "[email protected]"),
("Esp Classroom (MIT)", "[email protected]"),
]
create_users(
mit_realm, testsuite_mit_users, tos_version=settings.TERMS_OF_SERVICE_VERSION
)
mit_user = get_user_by_delivery_email("[email protected]", mit_realm)
bulk_create_streams(
mit_realm,
{
"core team": {
"description": "A private channel for core team members",
"invite_only": True,
}
},
)
core_team_stream = Stream.objects.get(name="core team", realm=mit_realm)
bulk_add_subscriptions(mit_realm, [core_team_stream], [mit_user], acting_user=None)
testsuite_lear_users = [
("King Lear", "[email protected]"),
("Cordelia, Lear's daughter", "[email protected]"),
]
create_users(