populate_db: Add alert words for Zulip users.

zerver/tests/test_alert_words.py

@@ -4,7 +4,7 @@ from zerver.lib.actions import do_add_alert_words, do_remove_alert_words
 from zerver.lib.alert_words import alert_words_in_realm, user_alert_words
 from zerver.lib.test_classes import ZulipTestCase
 from zerver.lib.test_helpers import most_recent_message, most_recent_usermessage
-from zerver.models import UserProfile
+from zerver.models import AlertWord, UserProfile
 
 
 class AlertWordTests(ZulipTestCase):
@@ -13,7 +13,11 @@ class AlertWordTests(ZulipTestCase):
     def get_user(self) -> UserProfile:
         # One nice thing about Hamlet is that he is
         # already subscribed to Denmark.
-        return self.example_user('hamlet')
+        user = self.example_user('hamlet')
+
+        # delete words from populate_db to simplify tests
+        AlertWord.objects.filter(user_profile=user).delete()
+        return user
 
     def test_internal_endpoint(self) -> None:
         user = self.get_user()
@@ -91,6 +95,10 @@ class AlertWordTests(ZulipTestCase):
         alert_words_in_realm. Alerts added for one user do not impact other
         users.
         """
+
+        # Clear all the words that we got from populate_db.
+        AlertWord.objects.all().delete()
+
         user1 = self.get_user()
 
         do_add_alert_words(user1, self.interesting_alert_word_list)

zerver/tests/test_message_flags.py

@@ -30,6 +30,16 @@ from zerver.models import (
 )
 
 
+def check_flags(flags: List[str], expected: Set[str]) -> None:
+    '''
+    The has_alert_word flag can be ignored for most tests.
+    '''
+    assert 'has_alert_word' not in expected
+    flag_set = set(flags)
+    flag_set.discard('has_alert_word')
+    if flag_set != expected:
+        raise AssertionError(f'expected flags (ignoring has_alert_word) to be {expected}')
+
 class FirstUnreadAnchorTests(ZulipTestCase):
     '''
     HISTORICAL NOTE:
@@ -167,7 +177,7 @@ class UnreadCountTests(ZulipTestCase):
         found = 0
         for msg in self.get_messages():
             if msg['id'] in self.unread_msg_ids:
-                self.assertEqual(msg['flags'], ['read'])
+                check_flags(msg['flags'], {'read'})
                 found += 1
         self.assertEqual(found, 2)
 
@@ -179,9 +189,9 @@ class UnreadCountTests(ZulipTestCase):
         # Ensure we properly remove just one flag
         for msg in self.get_messages():
             if msg['id'] == self.unread_msg_ids[0]:
-                self.assertEqual(msg['flags'], ['read'])
+                check_flags(msg['flags'], {'read'})
             elif msg['id'] == self.unread_msg_ids[1]:
-                self.assertEqual(msg['flags'], [])
+                check_flags(msg['flags'], set())
 
     def test_mark_all_in_stream_read(self) -> None:
         self.login('hamlet')
@@ -914,17 +924,17 @@ class MessageAccessTests(ZulipTestCase):
 
         for msg in self.get_messages():
             if msg['id'] in message_ids:
-                self.assertEqual(msg['flags'], ['starred'])
+                check_flags(msg['flags'], {'starred'})
             else:
-                self.assertEqual(msg['flags'], ['read'])
+                check_flags(msg['flags'], {'read'})
 
+        # Remove the stars.
         result = self.change_star(message_ids, False)
         self.assert_json_success(result)
 
-        # Remove the stars.
         for msg in self.get_messages():
             if msg['id'] in message_ids:
-                self.assertEqual(msg['flags'], [])
+                check_flags(msg['flags'], set())
 
     def test_change_star_public_stream_historical(self) -> None:
         """
@@ -981,11 +991,11 @@ class MessageAccessTests(ZulipTestCase):
 
         for msg in self.get_messages():
             if msg['id'] in message_ids:
-                self.assertEqual(set(msg['flags']), {'starred', 'historical', 'read'})
+                check_flags(msg['flags'], {'starred', 'historical', 'read'})
             elif msg['id'] in received_message_ids:
-                self.assertEqual(msg['flags'], [])
+                check_flags(msg['flags'], set())
             else:
-                self.assertEqual(msg['flags'], ['read'])
+                check_flags(msg['flags'], {'read'})
             self.assertNotIn(msg['id'], other_message_ids)
 
         result = self.change_star(message_ids, False)

zilencer/management/commands/populate_db.py

@@ -37,6 +37,7 @@ from zerver.lib.user_groups import create_user_group
 from zerver.lib.users import add_service
 from zerver.lib.utils import generate_api_key
 from zerver.models import (
+    AlertWord,
     Client,
     CustomProfileField,
     DefaultStream,
@@ -123,6 +124,41 @@ def subscribe_users_to_streams(realm: Realm, stream_dict: Dict[str, Dict[str, An
     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)
+        for i in range(4):
+            recs.append(
+                AlertWord(
+                    realm_id=realm_id,
+                    user_profile_id=user_id,
+                    word = alert_words[i],
+                )
+            )
+
+    AlertWord.objects.bulk_create(recs)
+
 class Command(BaseCommand):
     help = "Populate a test database"
 
@@ -522,6 +558,8 @@ class Command(BaseCommand):
         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()