notification_data: Add get_user_group_mentions_data function.

We will use this later to display which user group was mentioned
in push and email notifications.

`mentioned_user_group_ids` is kept as a List (not Set) to ensure proper
test coverage of the function, since it depends on the order of iteration,
and we cannot change the order of iteration for a set (which we'll need
to do for proper testing).

Part of #13080.

zerver/lib/notification_data.py

@@ -1,5 +1,7 @@
 from dataclasses import dataclass
-from typing import Collection, Optional, Set
+from typing import Collection, Dict, List, Optional, Set
+
+from zerver.lib.mention import MentionData
 
 
 @dataclass
@@ -109,3 +111,33 @@ class UserMessageNotificationsData:
             return "stream_email_notify"
         else:
             return None
+
+
+def get_user_group_mentions_data(
+    mentioned_user_ids: Set[int], mentioned_user_group_ids: List[int], mention_data: MentionData
+) -> Dict[int, int]:
+    # Maps user_id -> mentioned user_group_id
+    mentioned_user_groups_map: Dict[int, int] = dict()
+
+    # Add members of the mentioned user groups into `mentions_user_ids`.
+    for group_id in mentioned_user_group_ids:
+        member_ids = mention_data.get_group_members(group_id)
+        for member_id in member_ids:
+            if member_id in mentioned_user_ids:
+                # If a user is also mentioned personally, we use that as a trigger
+                # for notifications.
+                continue
+
+            if member_id in mentioned_user_groups_map:
+                # If multiple user groups are mentioned, we prefer the
+                # user group with the least members for email/mobile
+                # notifications.
+                previous_group_id = mentioned_user_groups_map[member_id]
+                previous_group_member_ids = mention_data.get_group_members(previous_group_id)
+
+                if len(previous_group_member_ids) > len(member_ids):
+                    mentioned_user_groups_map[member_id] = group_id
+            else:
+                mentioned_user_groups_map[member_id] = group_id
+
+    return mentioned_user_groups_map

zerver/tests/test_notification_data.py

@@ -1,4 +1,7 @@
+from zerver.lib.mention import MentionData
+from zerver.lib.notification_data import get_user_group_mentions_data
 from zerver.lib.test_classes import ZulipTestCase
+from zerver.lib.user_groups import create_user_group
 
 
 class TestNotificationData(ZulipTestCase):
@@ -308,3 +311,77 @@ class TestNotificationData(ZulipTestCase):
         self.assertTrue(
             user_data.is_notifiable(private_message=True, acting_user_id=acting_user_id, idle=True)
         )
+
+    def test_user_group_mentions_map(self) -> None:
+        hamlet = self.example_user("hamlet")
+        cordelia = self.example_user("cordelia")
+        realm = hamlet.realm
+
+        hamlet_only = create_user_group("hamlet_only", [hamlet], realm)
+        hamlet_and_cordelia = create_user_group("hamlet_and_cordelia", [hamlet, cordelia], realm)
+
+        # Base case. No user/user group mentions
+        result = get_user_group_mentions_data(
+            mentioned_user_ids=set(),
+            mentioned_user_group_ids=[],
+            mention_data=MentionData(realm.id, "no group mentioned"),
+        )
+        self.assertDictEqual(result, {})
+
+        # Only user group mentions, no personal mentions
+        result = get_user_group_mentions_data(
+            mentioned_user_ids=set(),
+            mentioned_user_group_ids=[hamlet_and_cordelia.id],
+            mention_data=MentionData(realm.id, "hey @*hamlet_and_cordelia*!"),
+        )
+        self.assertDictEqual(
+            result,
+            {
+                hamlet.id: hamlet_and_cordelia.id,
+                cordelia.id: hamlet_and_cordelia.id,
+            },
+        )
+
+        # Hamlet is mentioned in two user groups
+        # Test that we consider the smaller user group
+        result = get_user_group_mentions_data(
+            mentioned_user_ids=set(),
+            mentioned_user_group_ids=[hamlet_and_cordelia.id, hamlet_only.id],
+            mention_data=MentionData(realm.id, "hey @*hamlet_and_cordelia* and @*hamlet_only*"),
+        )
+        self.assertDictEqual(
+            result,
+            {
+                hamlet.id: hamlet_only.id,
+                cordelia.id: hamlet_and_cordelia.id,
+            },
+        )
+
+        # To make sure we aren't getting the expected data from over-writing in a loop,
+        # test the same setup as above, but with reversed group ids.
+        result = get_user_group_mentions_data(
+            mentioned_user_ids=set(),
+            mentioned_user_group_ids=[hamlet_only.id, hamlet_and_cordelia.id],
+            mention_data=MentionData(realm.id, "hey @*hamlet_only* and @*hamlet_and_cordelia*"),
+        )
+        self.assertDictEqual(
+            result,
+            {
+                hamlet.id: hamlet_only.id,
+                cordelia.id: hamlet_and_cordelia.id,
+            },
+        )
+
+        # Personal and user group mentioned. Test that we don't consider the user
+        # group mention for Hamlet in this case.
+        result = get_user_group_mentions_data(
+            mentioned_user_ids=set([hamlet.id]),
+            mentioned_user_group_ids=[hamlet_and_cordelia.id],
+            mention_data=MentionData(realm.id, "hey @*hamlet_and_cordelia*!"),
+        )
+        self.assertDictEqual(
+            result,
+            {
+                cordelia.id: hamlet_and_cordelia.id,
+            },
+        )