message_flags: Fix message flags not updating for wildcard mentions.

Earlier, editing a message with wildcard mention and removing the
wildcard mention didn't properly remove the corresponding flag for
it in the message object. It was only updating flags when mentions
were present in new message but not the other way around.

This commit fixes this behaviour by removing flags if the new message
removed mention from it.

zerver/actions/message_edit.py

@@ -406,6 +406,7 @@ def update_user_message_flags(
 ) -> None:
     mentioned_ids = rendering_result.mentions_user_ids
     ids_with_alert_words = rendering_result.user_ids_with_alert_words
+    stream_wildcard_mentioned = rendering_result.mentions_stream_wildcard
     changed_ums: set[UserMessage] = set()
 
     def update_flag(um: UserMessage, should_set: bool, flag: int) -> None:
@@ -425,11 +426,10 @@ def update_user_message_flags(
         mentioned = um.user_profile_id in mentioned_ids
         update_flag(um, mentioned, UserMessage.flags.mentioned)
 
-        if rendering_result.mentions_stream_wildcard:
-            update_flag(um, True, UserMessage.flags.stream_wildcard_mentioned)
-        elif rendering_result.mentions_topic_wildcard:
-            topic_wildcard_mentioned = um.user_profile_id in topic_participant_user_ids
-            update_flag(um, topic_wildcard_mentioned, UserMessage.flags.topic_wildcard_mentioned)
+        update_flag(um, stream_wildcard_mentioned, UserMessage.flags.stream_wildcard_mentioned)
+
+        topic_wildcard_mentioned = um.user_profile_id in topic_participant_user_ids
+        update_flag(um, topic_wildcard_mentioned, UserMessage.flags.topic_wildcard_mentioned)
 
     for um in changed_ums:
         um.save(update_fields=["flags"])

zerver/tests/test_message_edit.py

@@ -97,6 +97,20 @@ class EditMessageTest(ZulipTestCase):
                 message_edit_history,
             )
 
+    def check_message_flags(
+        self, message_id: int, user_ids: list[int], flag: str, check_present: bool
+    ) -> None:
+        # Make sure we updated the message flags correctly to the DB.
+        for user_id in user_ids:
+            um = UserMessage.objects.get(
+                user_profile_id=user_id,
+                message_id=message_id,
+            )
+            if check_present:
+                self.assertIn(flag, um.flags_list())
+            else:
+                self.assertNotIn(flag, um.flags_list())
+
     def test_edit_message_no_changes(self) -> None:
         self.login("hamlet")
         msg_id = self.send_stream_message(
@@ -1693,6 +1707,13 @@ class EditMessageTest(ZulipTestCase):
         self.login_user(hamlet)
         message_id = self.send_stream_message(hamlet, stream_name, "Hello everyone")
 
+        self.check_message_flags(
+            message_id,
+            user_ids=[hamlet.id, cordelia.id],
+            flag="topic_wildcard_mentioned",
+            check_present=False,
+        )
+
         users_to_be_notified = sorted(
             [
                 {
@@ -1728,6 +1749,74 @@ class EditMessageTest(ZulipTestCase):
                 called = True
         self.assertTrue(called)
 
+        self.check_message_flags(
+            message_id, user_ids=[hamlet.id], flag="topic_wildcard_mentioned", check_present=True
+        )
+        self.check_message_flags(
+            message_id, user_ids=[cordelia.id], flag="topic_wildcard_mentioned", check_present=False
+        )
+
+    @mock.patch("zerver.actions.message_edit.send_event_on_commit")
+    def test_remove_topic_wildcard_mention(self, mock_send_event: mock.MagicMock) -> None:
+        stream_name = "Macbeth"
+        hamlet = self.example_user("hamlet")
+        cordelia = self.example_user("cordelia")
+        self.make_stream(stream_name, history_public_to_subscribers=True)
+        self.subscribe(hamlet, stream_name)
+        self.subscribe(cordelia, stream_name)
+        self.login_user(hamlet)
+        message_id = self.send_stream_message(hamlet, stream_name, "Hello @**topic**")
+
+        self.check_message_flags(
+            message_id, user_ids=[hamlet.id], flag="topic_wildcard_mentioned", check_present=True
+        )
+        self.check_message_flags(
+            message_id, user_ids=[cordelia.id], flag="topic_wildcard_mentioned", check_present=False
+        )
+
+        users_to_be_notified = sorted(
+            [
+                {
+                    "id": hamlet.id,
+                    "flags": ["read"],
+                },
+                {
+                    "id": cordelia.id,
+                    "flags": [],
+                },
+            ],
+            key=itemgetter("id"),
+        )
+        result = self.client_patch(
+            f"/json/messages/{message_id}",
+            {
+                "content": "Hello everyone",
+            },
+        )
+        self.assert_json_success(result)
+
+        # Extract the send_event call where event type is 'update_message'.
+        # Here we assert 'stream_wildcard_mention_user_ids' is empty and flag
+        # is removed.
+        called = False
+        for call_args in mock_send_event.call_args_list:
+            (arg_realm, arg_event, arg_notified_users) = call_args[0]
+            if arg_event["type"] == "update_message":
+                self.assertEqual(arg_event["type"], "update_message")
+                self.assertEqual(arg_event["stream_wildcard_mention_user_ids"], [])
+                self.assertEqual(
+                    sorted(arg_notified_users, key=itemgetter("id")), users_to_be_notified
+                )
+                called = True
+        self.assertTrue(called)
+
+        self.check_message_flags(
+            message_id,
+            user_ids=[hamlet.id, cordelia.id],
+            flag="topic_wildcard_mentioned",
+            check_present=False,
+        )
+
     def test_topic_wildcard_mention_restrictions_when_editing(self) -> None:
         cordelia = self.example_user("cordelia")
         shiva = self.example_user("shiva")
@@ -1804,6 +1893,13 @@ class EditMessageTest(ZulipTestCase):
         self.login_user(hamlet)
         message_id = self.send_stream_message(hamlet, stream_name, "Hello everyone")
 
+        self.check_message_flags(
+            message_id,
+            user_ids=[hamlet.id, cordelia.id],
+            flag="stream_wildcard_mentioned",
+            check_present=False,
+        )
+
         users_to_be_notified = sorted(
             [
                 {
@@ -1841,6 +1937,74 @@ class EditMessageTest(ZulipTestCase):
                 called = True
         self.assertTrue(called)
 
+        self.check_message_flags(
+            message_id,
+            user_ids=[hamlet.id, cordelia.id],
+            flag="stream_wildcard_mentioned",
+            check_present=True,
+        )
+
+    @mock.patch("zerver.actions.message_edit.send_event_on_commit")
+    def test_remove_stream_wildcard_mention(self, mock_send_event: mock.MagicMock) -> None:
+        stream_name = "Macbeth"
+        hamlet = self.example_user("hamlet")
+        cordelia = self.example_user("cordelia")
+        self.make_stream(stream_name, history_public_to_subscribers=True)
+        self.subscribe(hamlet, stream_name)
+        self.subscribe(cordelia, stream_name)
+        self.login_user(hamlet)
+        message_id = self.send_stream_message(hamlet, stream_name, "Hello @**everyone**")
+
+        self.check_message_flags(
+            message_id,
+            user_ids=[hamlet.id, cordelia.id],
+            flag="stream_wildcard_mentioned",
+            check_present=True,
+        )
+
+        users_to_be_notified = sorted(
+            [
+                {
+                    "id": hamlet.id,
+                    "flags": ["read"],
+                },
+                {
+                    "id": cordelia.id,
+                    "flags": [],
+                },
+            ],
+            key=itemgetter("id"),
+        )
+        result = self.client_patch(
+            f"/json/messages/{message_id}",
+            {
+                "content": "Hello everyone",
+            },
+        )
+        self.assert_json_success(result)
+
+        # Extract the send_event call where event type is 'update_message'.
+        # Here we assert 'stream_wildcard_mention_user_ids' is empty and flag
+        # is removed.
+        called = False
+        for call_args in mock_send_event.call_args_list:
+            (arg_realm, arg_event, arg_notified_users) = call_args[0]
+            if arg_event["type"] == "update_message":
+                self.assertEqual(arg_event["type"], "update_message")
+                self.assertEqual(arg_event["stream_wildcard_mention_user_ids"], [])
+                self.assertEqual(
+                    sorted(arg_notified_users, key=itemgetter("id")), users_to_be_notified
+                )
+                called = True
+        self.assertTrue(called)
+
+        self.check_message_flags(
+            message_id,
+            user_ids=[hamlet.id, cordelia.id],
+            flag="stream_wildcard_mentioned",
+            check_present=False,
+        )
+
     def test_stream_wildcard_mention_restrictions_when_editing(self) -> None:
         cordelia = self.example_user("cordelia")
         shiva = self.example_user("shiva")