ref(notifications): Refactor notification utilities for msteams (#36883)

* Move notification utility functions to a common file for reuse in msteams notifications.
* Make a slight modification to the Identity query to exclude msteams identity rows which have a `team_id` as `external_id` instead of a `user_id`.

src/sentry/integrations/notifications.py

@@ -0,0 +1,98 @@
+from __future__ import annotations
+
+from collections import defaultdict
+from typing import Any, Iterable, Mapping, MutableMapping
+
+from django.db.models import F
+
+from sentry.constants import ObjectStatus
+from sentry.models import ExternalActor, Identity, Integration, Organization, Team, User
+from sentry.notifications.notifications.base import BaseNotification
+from sentry.types.integrations import EXTERNAL_PROVIDERS, ExternalProviders
+
+
+def get_context(
+    notification: BaseNotification,
+    recipient: Team | User,
+    shared_context: Mapping[str, Any],
+    extra_context: Mapping[str, Any],
+) -> Mapping[str, Any]:
+    """Compose the various levels of context and add Slack-specific fields."""
+    return {
+        **shared_context,
+        **notification.get_recipient_context(recipient, extra_context),
+    }
+
+
+def get_channel_and_integration_by_user(
+    user: User,
+    organization: Organization,
+    provider: ExternalProviders,
+) -> Mapping[str, Integration]:
+
+    identities = (
+        Identity.objects.filter(
+            idp__type=EXTERNAL_PROVIDERS[provider],
+            user=user.id,
+        )
+        # For Microsoft Teams integration, initially we create rows in the
+        # identity table with the external_id as a team_id instead of the user_id.
+        # We need to exclude rows where this is NOT updated to the user_id later.
+        .exclude(external_id=F("idp__external_id")).select_related("idp")
+    )
+
+    if not identities:
+        # The user may not have linked their identity so just move on
+        # since there are likely other users or teams in the list of
+        # recipients.
+        return {}
+
+    integrations = Integration.objects.get_active_integrations(organization.id).filter(
+        provider=EXTERNAL_PROVIDERS[provider],
+        external_id__in=[identity.idp.external_id for identity in identities],
+    )
+
+    channels_to_integration = {}
+    for identity in identities:
+        for integration in integrations:
+            if identity.idp.external_id == integration.external_id:
+                channels_to_integration[identity.external_id] = integration
+                break
+
+    return channels_to_integration
+
+
+def get_channel_and_integration_by_team(
+    team: Team, organization: Organization, provider: ExternalProviders
+) -> Mapping[str, Integration]:
+    try:
+        external_actor = (
+            ExternalActor.objects.filter(
+                provider=provider.value,
+                actor_id=team.actor_id,
+                organization=organization,
+                integration__status=ObjectStatus.ACTIVE,
+                integration__organizationintegration__status=ObjectStatus.ACTIVE,
+                # limit to org here to prevent multiple query results
+                integration__organizationintegration__organization=organization,
+            )
+            .select_related("integration")
+            .get()
+        )
+    except ExternalActor.DoesNotExist:
+        return {}
+    return {external_actor.external_id: external_actor.integration}
+
+
+def get_integrations_by_channel_by_recipient(
+    organization: Organization, recipients: Iterable[Team | User], provider: ExternalProviders
+) -> MutableMapping[Team | User, Mapping[str, Integration]]:
+    output: MutableMapping[Team | User, Mapping[str, Integration]] = defaultdict(dict)
+    for recipient in recipients:
+        channels_to_integrations = (
+            get_channel_and_integration_by_user(recipient, organization, provider)
+            if isinstance(recipient, User)
+            else get_channel_and_integration_by_team(recipient, organization, provider)
+        )
+        output[recipient] = channels_to_integrations
+    return output

src/sentry/integrations/slack/notifications.py

@@ -1,24 +1,23 @@
 from __future__ import annotations
 
 import logging
-from collections import defaultdict
 from copy import copy
-from typing import Any, Iterable, List, Mapping, MutableMapping
+from typing import Any, Iterable, List, Mapping
 
 import sentry_sdk
 
-from sentry.constants import ObjectStatus
 from sentry.integrations.mixins import NotifyBasicMixin
+from sentry.integrations.notifications import get_context, get_integrations_by_channel_by_recipient
 from sentry.integrations.slack.client import SlackClient
 from sentry.integrations.slack.message_builder import SlackAttachment
 from sentry.integrations.slack.message_builder.notifications import get_message_builder
-from sentry.models import ExternalActor, Identity, Integration, Organization, Team, User
+from sentry.models import Integration, Team, User
 from sentry.notifications.additional_attachment_manager import get_additional_attachment
 from sentry.notifications.notifications.base import BaseNotification
 from sentry.notifications.notify import register_notification_provider
 from sentry.shared_integrations.exceptions import ApiError
 from sentry.tasks.integrations.slack import post_message
-from sentry.types.integrations import EXTERNAL_PROVIDERS, ExternalProviders
+from sentry.types.integrations import ExternalProviders
 from sentry.utils import json, metrics
 
 logger = logging.getLogger("sentry.notifications")
@@ -59,85 +58,6 @@ def get_attachments(
     return [attachments]
 
 
-def get_context(
-    notification: BaseNotification,
-    recipient: Team | User,
-    shared_context: Mapping[str, Any],
-    extra_context: Mapping[str, Any],
-) -> Mapping[str, Any]:
-    """Compose the various levels of context and add Slack-specific fields."""
-    return {
-        **shared_context,
-        **notification.get_recipient_context(recipient, extra_context),
-    }
-
-
-def get_channel_and_integration_by_user(
-    user: User, organization: Organization
-) -> Mapping[str, Integration]:
-
-    identities = Identity.objects.filter(
-        idp__type=EXTERNAL_PROVIDERS[ExternalProviders.SLACK],
-        user=user.id,
-    ).select_related("idp")
-
-    if not identities:
-        # The user may not have linked their identity so just move on
-        # since there are likely other users or teams in the list of
-        # recipients.
-        return {}
-
-    integrations = Integration.objects.get_active_integrations(organization.id).filter(
-        provider=EXTERNAL_PROVIDERS[ExternalProviders.SLACK],
-        external_id__in=[identity.idp.external_id for identity in identities],
-    )
-
-    channels_to_integration = {}
-    for identity in identities:
-        for integration in integrations:
-            if identity.idp.external_id == integration.external_id:
-                channels_to_integration[identity.external_id] = integration
-                break
-
-    return channels_to_integration
-
-
-def get_channel_and_integration_by_team(
-    team: Team, organization: Organization
-) -> Mapping[str, Integration]:
-    try:
-        external_actor = (
-            ExternalActor.objects.filter(
-                provider=ExternalProviders.SLACK.value,
-                actor_id=team.actor_id,
-                organization=organization,
-                integration__status=ObjectStatus.ACTIVE,
-                integration__organizationintegration__status=ObjectStatus.ACTIVE,
-                # limit to org here to prevent multiple query results
-                integration__organizationintegration__organization=organization,
-            )
-            .select_related("integration")
-            .get()
-        )
-    except ExternalActor.DoesNotExist:
-        return {}
-    return {external_actor.external_id: external_actor.integration}
-
-
-def get_integrations_by_channel_by_recipient(
-    organization: Organization, recipients: Iterable[Team | User]
-) -> MutableMapping[Team | User, Mapping[str, Integration]]:
-    output: MutableMapping[Team | User, Mapping[str, Integration]] = defaultdict(dict)
-    for recipient in recipients:
-        channels_to_integrations = (
-            get_channel_and_integration_by_user(recipient, organization)
-            if isinstance(recipient, User)
-            else get_channel_and_integration_by_team(recipient, organization)
-        )
-        output[recipient] = channels_to_integrations
-    return output
-
-
 def _notify_recipient(
     notification: BaseNotification,
     recipient: Team | User,
@@ -196,7 +116,9 @@ def send_notification_as_slack(
     with sentry_sdk.start_span(
         op="notification.send_slack", description="gen_channel_integration_map"
     ):
-        data = get_integrations_by_channel_by_recipient(notification.organization, recipients)
+        data = get_integrations_by_channel_by_recipient(
+            notification.organization, recipients, ExternalProviders.SLACK
+        )
     for recipient, integrations_by_channel in data.items():
         with sentry_sdk.start_span(op="notification.send_slack", description="send_one"):
             with sentry_sdk.start_span(op="notification.send_slack", description="gen_attachments"):

src/sentry/testutils/fixtures.py

@@ -11,6 +11,7 @@ from sentry.models import (
     OrganizationMember,
     OrganizationMemberTeam,
 )
+from sentry.models.user import User
 from sentry.testutils.factories import Factories
 from sentry.testutils.helpers.datetime import before_now, iso_format
 
@@ -375,13 +376,18 @@ class Fixtures:
         self,
         organization: "Organization",
         external_id: str = "TXXXXXXX1",
+        user: User = None,
+        identity_external_id: str = "UXXXXXXX1",
         **kwargs: Any,
     ):
+        if user is None:
+            user = organization.get_default_owner()
+
         integration = Factories.create_slack_integration(
             organization=organization, external_id=external_id, **kwargs
         )
         idp = Factories.create_identity_provider(integration=integration)
-        Factories.create_identity(organization.get_default_owner(), idp, "UXXXXXXX1")
+        Factories.create_identity(user, idp, identity_external_id)
 
         return integration
 

tests/sentry/integrations/test_notification_utilities.py

@@ -0,0 +1,57 @@
+from sentry.integrations.notifications import get_integrations_by_channel_by_recipient
+from sentry.testutils.cases import TestCase
+from sentry.testutils.helpers.notifications import DummyNotification
+from sentry.types.integrations import ExternalProviders
+
+
+class TestNotificationUtilities(TestCase):
+    def setUp(self):
+        super().setUp()
+        self.notification = DummyNotification(self.organization)
+
+        self.external_user_id_1 = "UXXXXXXX1"
+        self.integration = self.create_slack_integration(self.notification.organization)
+
+        self.user_2 = self.create_user()
+        self.external_team_id_2 = "TXXXXXXX2"
+        self.integration2 = self.create_slack_integration(
+            self.notification.organization,
+            external_id=self.external_team_id_2,
+            user=self.user_2,
+            identity_external_id=self.external_team_id_2,
+        )
+
+    def test_simple(self):
+        integrations_by_channel_by_recipient = get_integrations_by_channel_by_recipient(
+            self.notification.organization,
+            [self.user],
+            ExternalProviders.SLACK,
+        )
+
+        assert {
+            self.user: {self.external_user_id_1: self.integration}
+        } == integrations_by_channel_by_recipient
+
+    def test_matching_idp_and_identity_external_id(self):
+        """
+        Test that rows where identity.external_id is equal to idp.external_id are excluded.
+        """
+        integrations_by_channel_by_recipient = get_integrations_by_channel_by_recipient(
+            self.notification.organization,
+            [self.user_2],
+            ExternalProviders.SLACK,
+        )
+
+        assert {self.user_2: {}} == integrations_by_channel_by_recipient
+
+    def test_multiple(self):
+        integrations_by_channel_by_recipient = get_integrations_by_channel_by_recipient(
+            self.notification.organization,
+            [self.user, self.user_2],
+            ExternalProviders.SLACK,
+        )
+
+        assert {
+            self.user: {self.external_user_id_1: self.integration},
+            self.user_2: {},
+        } == integrations_by_channel_by_recipient