test_e2ee_push_notification: Improve tests to cover more cases.

This commit adds a test and updates a few existing tests to
cover more cases related to send push notifications.

* We no longer mock the 'send_push_notifications_legacy' function
  while testing 'send_push_notifications' codepath and vice-versa.
  This makes the tests more realistic as both functions gets called
  in 'handle_push_notification'.

  This covers the case when only old clients (which don't support
  E2EE) exists for a user. Or only updated clients (which supports
  E2EE) exist.

* Adds a test 'test_both_old_and_new_client_coexists' for the case
  when a user has both type of clients at an instant i.e. they have
  updated a few devices only.

zerver/lib/push_notifications.py

@@ -1317,11 +1317,7 @@ def send_push_notifications_legacy(
         PushDeviceToken.objects.filter(user=user_profile, kind=PushDeviceToken.APNS).order_by("id")
     )
 
-    # Added this nocoverage check for this commit. We plan to add a commit
-    # which will remove the mocking of 'send_push_notifications_legacy'
-    # in 'test_e2ee_push_notifications' - and this case will be covered
-    # there as a part of that broader effort.
-    if len(android_devices) + len(apple_devices) == 0:  # nocoverage
+    if len(android_devices) + len(apple_devices) == 0:
         logger.info(
             "Skipping legacy push notifications for user %s because there are no registered devices",
             user_profile.id,

zerver/lib/test_classes.py

@@ -2893,30 +2893,71 @@ class E2EEPushNotificationTestCase(BouncerTestCase):
 
         return registered_device_apple, registered_device_android
 
-    @contextmanager
-    def mock_fcm(self) -> Iterator[mock.MagicMock]:
-        with mock.patch("zilencer.lib.push_notifications.firebase_messaging") as mock_fcm_messaging:
-            yield mock_fcm_messaging
+    def register_old_push_devices_for_notification(self) -> tuple[PushDeviceToken, PushDeviceToken]:
+        hamlet = self.example_user("hamlet")
+
+        registered_device_android = PushDeviceToken.objects.create(
+            kind=PushDeviceToken.FCM,
+            token="token-fcm",
+            user=hamlet,
+            ios_app_id=None,
+        )
+        registered_device_apple = PushDeviceToken.objects.create(
+            kind=PushDeviceToken.APNS,
+            token="token-apns",
+            user=hamlet,
+            ios_app_id="abc",
+        )
+        return registered_device_apple, registered_device_android
 
     @contextmanager
-    def mock_apns(self) -> Iterator[mock.AsyncMock]:
+    def mock_fcm(self, for_legacy: bool = False) -> Iterator[mock.MagicMock]:
+        if for_legacy:
+            with (
+                mock.patch("zerver.lib.push_notifications.fcm_app"),
+                mock.patch(
+                    "zerver.lib.push_notifications.firebase_messaging"
+                ) as mock_fcm_messaging,
+            ):
+                yield mock_fcm_messaging
+        else:
+            with (
+                mock.patch("zilencer.lib.push_notifications.fcm_app"),
+                mock.patch(
+                    "zilencer.lib.push_notifications.firebase_messaging"
+                ) as mock_fcm_messaging,
+            ):
+                yield mock_fcm_messaging
+
+    @contextmanager
+    def mock_apns(self, for_legacy: bool = False) -> Iterator[mock.AsyncMock]:
         apns = mock.Mock(spec=aioapns.APNs)
         apns.send_notification = mock.AsyncMock()
         apns_context = APNsContext(
             apns=apns,
             loop=asyncio.new_event_loop(),
         )
+        target = (
+            "zerver.lib.push_notifications.get_apns_context"
+            if for_legacy
+            else "zilencer.lib.push_notifications.get_apns_context"
+        )
         try:
-            with mock.patch("zilencer.lib.push_notifications.get_apns_context") as mock_get:
+            with mock.patch(target) as mock_get:
                 mock_get.return_value = apns_context
                 yield apns.send_notification
         finally:
             apns_context.loop.close()
 
-    def make_fcm_success_response(self) -> firebase_messaging.BatchResponse:
-        device_ids_count = RemotePushDevice.objects.filter(
-            token_kind=RemotePushDevice.TokenKind.FCM
-        ).count()
+    def make_fcm_success_response(
+        self, for_legacy: bool = False
+    ) -> firebase_messaging.BatchResponse:
+        if for_legacy:
+            device_ids_count = PushDeviceToken.objects.filter(kind=PushDeviceToken.FCM).count()
+        else:
+            device_ids_count = RemotePushDevice.objects.filter(
+                token_kind=RemotePushDevice.TokenKind.FCM
+            ).count()
         responses = [
             firebase_messaging.SendResponse(exception=None, resp=dict(name=str(idx)))
             for idx in range(device_ids_count)

zerver/tests/test_e2ee_push_notifications.py

@@ -16,9 +16,8 @@ from zilencer.models import RemoteRealm, RemoteRealmCount
 
 
 @activate_push_notification_service()
-@mock.patch("zerver.lib.push_notifications.send_push_notifications_legacy")
 class SendPushNotificationTest(E2EEPushNotificationTestCase):
-    def test_success_cloud(self, unused_mock: mock.MagicMock) -> None:
+    def test_success_cloud(self) -> None:
         hamlet = self.example_user("hamlet")
         aaron = self.example_user("aaron")
 
@@ -56,9 +55,14 @@ class SendPushNotificationTest(E2EEPushNotificationTestCase):
             )
             self.assertEqual(
                 "INFO:zerver.lib.push_notifications:"
-                f"APNs: Success sending to (push_account_id={registered_device_apple.push_account_id}, device={registered_device_apple.token})",
+                f"Skipping legacy push notifications for user {hamlet.id} because there are no registered devices",
                 zerver_logger.output[1],
             )
+            self.assertEqual(
+                "INFO:zerver.lib.push_notifications:"
+                f"APNs: Success sending to (push_account_id={registered_device_apple.push_account_id}, device={registered_device_apple.token})",
+                zerver_logger.output[2],
+            )
             self.assertEqual(
                 "INFO:zilencer.lib.push_notifications:"
                 f"FCM: Sent message with ID: 0 to (push_account_id={registered_device_android.push_account_id}, device={registered_device_android.token})",
@@ -67,7 +71,7 @@ class SendPushNotificationTest(E2EEPushNotificationTestCase):
             self.assertEqual(
                 "INFO:zerver.lib.push_notifications:"
                 f"Sent E2EE mobile push notifications for user {hamlet.id}: 1 via FCM, 1 via APNs",
-                zerver_logger.output[2],
+                zerver_logger.output[3],
             )
 
             realm_count_dict = (
@@ -77,7 +81,7 @@ class SendPushNotificationTest(E2EEPushNotificationTestCase):
             )
             self.assertEqual(realm_count_dict, dict(subgroup=None, value=2))
 
-    def test_no_registered_device(self, unused_mock: mock.MagicMock) -> None:
+    def test_no_registered_device(self) -> None:
         aaron = self.example_user("aaron")
         hamlet = self.example_user("hamlet")
 
@@ -99,11 +103,16 @@ class SendPushNotificationTest(E2EEPushNotificationTestCase):
             )
             self.assertEqual(
                 "INFO:zerver.lib.push_notifications:"
-                f"Skipping E2EE push notifications for user {hamlet.id} because there are no registered devices",
+                f"Skipping legacy push notifications for user {hamlet.id} because there are no registered devices",
                 zerver_logger.output[1],
             )
+            self.assertEqual(
+                "INFO:zerver.lib.push_notifications:"
+                f"Skipping E2EE push notifications for user {hamlet.id} because there are no registered devices",
+                zerver_logger.output[2],
+            )
 
-    def test_invalid_or_expired_token(self, unused_mock: mock.MagicMock) -> None:
+    def test_invalid_or_expired_token(self) -> None:
         aaron = self.example_user("aaron")
         hamlet = self.example_user("hamlet")
 
@@ -147,7 +156,7 @@ class SendPushNotificationTest(E2EEPushNotificationTestCase):
             self.assertEqual(
                 "INFO:zerver.lib.push_notifications:"
                 f"APNs: Removing invalid/expired token {registered_device_apple.token} (BadDeviceToken)",
-                zerver_logger.output[1],
+                zerver_logger.output[2],
             )
             self.assertEqual(
                 "INFO:zilencer.lib.push_notifications:"
@@ -157,12 +166,12 @@ class SendPushNotificationTest(E2EEPushNotificationTestCase):
             self.assertEqual(
                 "INFO:zerver.lib.push_notifications:"
                 f"Deleting PushDevice rows with the following device IDs based on response from bouncer: [{registered_device_apple.device_id}, {registered_device_android.device_id}]",
-                zerver_logger.output[2],
+                zerver_logger.output[3],
             )
             self.assertEqual(
                 "INFO:zerver.lib.push_notifications:"
                 f"Sent E2EE mobile push notifications for user {hamlet.id}: 0 via FCM, 0 via APNs",
-                zerver_logger.output[3],
+                zerver_logger.output[4],
             )
 
             # Verify `expired_time` set for `RemotePushDevice` entries
@@ -173,7 +182,7 @@ class SendPushNotificationTest(E2EEPushNotificationTestCase):
             self.assertIsNotNone(registered_device_android.expired_time)
             self.assertEqual(PushDevice.objects.count(), 0)
 
-    def test_fcm_apns_error(self, unused_mock: mock.MagicMock) -> None:
+    def test_fcm_apns_error(self) -> None:
         hamlet = self.example_user("hamlet")
         aaron = self.example_user("aaron")
 
@@ -220,7 +229,7 @@ class SendPushNotificationTest(E2EEPushNotificationTestCase):
             self.assertEqual(
                 "INFO:zerver.lib.push_notifications:"
                 f"Sent E2EE mobile push notifications for user {hamlet.id}: 0 via FCM, 0 via APNs",
-                zerver_logger.output[1],
+                zerver_logger.output[2],
             )
 
         # `firebase_messaging.send_each` raises Error.
@@ -255,10 +264,10 @@ class SendPushNotificationTest(E2EEPushNotificationTestCase):
             self.assertEqual(
                 "INFO:zerver.lib.push_notifications:"
                 f"Sent E2EE mobile push notifications for user {hamlet.id}: 0 via FCM, 1 via APNs",
-                zerver_logger.output[1],
+                zerver_logger.output[2],
             )
 
-    def test_early_return_if_expired_time_set(self, unused_mock: mock.MagicMock) -> None:
+    def test_early_return_if_expired_time_set(self) -> None:
         aaron = self.example_user("aaron")
         hamlet = self.example_user("hamlet")
 
@@ -299,7 +308,7 @@ class SendPushNotificationTest(E2EEPushNotificationTestCase):
 
     @responses.activate
     @override_settings(ZILENCER_ENABLED=False)
-    def test_success_self_hosted(self, unused_mock: mock.MagicMock) -> None:
+    def test_success_self_hosted(self) -> None:
         self.add_mock_response()
 
         hamlet = self.example_user("hamlet")
@@ -352,9 +361,14 @@ class SendPushNotificationTest(E2EEPushNotificationTestCase):
             )
             self.assertEqual(
                 "INFO:zerver.lib.push_notifications:"
-                f"APNs: Success sending to (push_account_id={registered_device_apple.push_account_id}, device={registered_device_apple.token})",
+                f"Skipping legacy push notifications for user {hamlet.id} because there are no registered devices",
                 zerver_logger.output[1],
             )
+            self.assertEqual(
+                "INFO:zerver.lib.push_notifications:"
+                f"APNs: Success sending to (push_account_id={registered_device_apple.push_account_id}, device={registered_device_apple.token})",
+                zerver_logger.output[2],
+            )
             self.assertEqual(
                 "INFO:zilencer.lib.push_notifications:"
                 f"FCM: Sent message with ID: 0 to (push_account_id={registered_device_android.push_account_id}, device={registered_device_android.token})",
@@ -363,7 +377,7 @@ class SendPushNotificationTest(E2EEPushNotificationTestCase):
             self.assertEqual(
                 "INFO:zerver.lib.push_notifications:"
                 f"Sent E2EE mobile push notifications for user {hamlet.id}: 1 via FCM, 1 via APNs",
-                zerver_logger.output[2],
+                zerver_logger.output[3],
             )
 
             realm_count_dict = (
@@ -393,7 +407,7 @@ class SendPushNotificationTest(E2EEPushNotificationTestCase):
 
     @responses.activate
     @override_settings(ZILENCER_ENABLED=False)
-    def test_missing_remote_realm_error(self, unused_mock: mock.MagicMock) -> None:
+    def test_missing_remote_realm_error(self) -> None:
         self.add_mock_response()
 
         hamlet = self.example_user("hamlet")
@@ -438,7 +452,7 @@ class SendPushNotificationTest(E2EEPushNotificationTestCase):
             self.assertEqual(
                 "WARNING:zerver.lib.push_notifications:"
                 "Bouncer refused to send E2EE push notification: Organization not registered",
-                zerver_logger.output[1],
+                zerver_logger.output[2],
             )
 
             realm.refresh_from_db()
@@ -447,7 +461,7 @@ class SendPushNotificationTest(E2EEPushNotificationTestCase):
 
     @responses.activate
     @override_settings(ZILENCER_ENABLED=False)
-    def test_no_plan_error(self, unused_mock: mock.MagicMock) -> None:
+    def test_no_plan_error(self) -> None:
         self.add_mock_response()
 
         hamlet = self.example_user("hamlet")
@@ -488,18 +502,106 @@ class SendPushNotificationTest(E2EEPushNotificationTestCase):
                 "WARNING:zerver.lib.push_notifications:"
                 "Bouncer refused to send E2EE push notification: Your plan doesn't allow sending push notifications. "
                 "Reason provided by the server: Push notifications access with 10+ users requires signing up for a plan. https://zulip.com/plans/",
-                zerver_logger.output[1],
+                zerver_logger.output[2],
             )
 
             realm.refresh_from_db()
             self.assertFalse(realm.push_notifications_enabled)
             self.assertIsNone(realm.push_notifications_enabled_end_timestamp)
 
+    def test_both_old_and_new_client_coexists(self) -> None:
+        """Test coexistence of old (which don't support E2EE)
+        and new client devices registered for push notifications.
+        """
+        hamlet = self.example_user("hamlet")
+        aaron = self.example_user("aaron")
+
+        registered_device_apple, registered_device_android = (
+            self.register_push_devices_for_notification()
+        )
+        registered_device_apple_old, registered_device_android_old = (
+            self.register_old_push_devices_for_notification()
+        )
+
+        message_id = self.send_personal_message(
+            from_user=aaron, to_user=hamlet, skip_capture_on_commit_callbacks=True
+        )
+        missed_message = {
+            "message_id": message_id,
+            "trigger": NotificationTriggers.DIRECT_MESSAGE,
+        }
+
+        self.assertEqual(RealmCount.objects.count(), 0)
+
+        with (
+            self.mock_fcm(for_legacy=True) as mock_fcm_messaging_legacy,
+            self.mock_apns(for_legacy=True) as send_notification_legacy,
+            self.mock_fcm() as mock_fcm_messaging,
+            self.mock_apns() as send_notification,
+            mock.patch(
+                "zerver.lib.push_notifications.uses_notification_bouncer", return_value=False
+            ),
+            self.assertLogs("zerver.lib.push_notifications", level="INFO") as zerver_logger,
+            self.assertLogs("zilencer.lib.push_notifications", level="INFO") as zilencer_logger,
+        ):
+            mock_fcm_messaging_legacy.send_each.return_value = self.make_fcm_success_response(
+                for_legacy=True
+            )
+            send_notification_legacy.return_value.is_successful = True
+            mock_fcm_messaging.send_each.return_value = self.make_fcm_success_response()
+            send_notification.return_value.is_successful = True
+
+            handle_push_notification(hamlet.id, missed_message)
+
+            mock_fcm_messaging_legacy.send_each.assert_called_once()
+            send_notification_legacy.assert_called_once()
+            mock_fcm_messaging.send_each.assert_called_once()
+            send_notification.assert_called_once()
+
+            self.assertEqual(
+                "INFO:zerver.lib.push_notifications:"
+                f"Sending push notifications to mobile clients for user {hamlet.id}",
+                zerver_logger.output[0],
+            )
+            # Logs in legacy codepath
+            self.assertEqual(
+                zerver_logger.output[1:6],
+                [
+                    f"INFO:zerver.lib.push_notifications:Sending mobile push notifications for local user {hamlet.id}: 1 via FCM devices, 1 via APNs devices",
+                    f"INFO:zerver.lib.push_notifications:APNs: Sending notification for local user <id:{hamlet.id}> to 1 devices (skipped 0 duplicates)",
+                    f"INFO:zerver.lib.push_notifications:APNs: Success sending for user <id:{hamlet.id}> to device {registered_device_apple_old.token}",
+                    f"INFO:zerver.lib.push_notifications:FCM: Sending notification for local user <id:{hamlet.id}> to 1 devices",
+                    f"INFO:zerver.lib.push_notifications:FCM: Sent message with ID: 0 to {registered_device_android_old.token}",
+                ],
+            )
+            # Logs in E2EE codepath
+            self.assertEqual(
+                "INFO:zerver.lib.push_notifications:"
+                f"APNs: Success sending to (push_account_id={registered_device_apple.push_account_id}, device={registered_device_apple.token})",
+                zerver_logger.output[6],
+            )
+            self.assertEqual(
+                "INFO:zilencer.lib.push_notifications:"
+                f"FCM: Sent message with ID: 0 to (push_account_id={registered_device_android.push_account_id}, device={registered_device_android.token})",
+                zilencer_logger.output[0],
+            )
+            self.assertEqual(
+                "INFO:zerver.lib.push_notifications:"
+                f"Sent E2EE mobile push notifications for user {hamlet.id}: 1 via FCM, 1 via APNs",
+                zerver_logger.output[7],
+            )
+
+            realm_count_dict = (
+                RealmCount.objects.filter(property="mobile_pushes_sent::day")
+                .values("subgroup", "value")
+                .last()
+            )
+            self.assertEqual(realm_count_dict, dict(subgroup=None, value=4))
+
 
 @activate_push_notification_service()
-@mock.patch("zerver.lib.push_notifications.send_push_notifications_legacy")
 class RemovePushNotificationTest(E2EEPushNotificationTestCase):
-    def test_success_cloud(self, unused_mock: mock.MagicMock) -> None:
+    def test_success_cloud(self) -> None:
         hamlet = self.example_user("hamlet")
         aaron = self.example_user("aaron")
 
@@ -531,12 +633,12 @@ class RemovePushNotificationTest(E2EEPushNotificationTestCase):
             self.assertEqual(
                 "INFO:zerver.lib.push_notifications:"
                 f"Sent E2EE mobile push notifications for user {hamlet.id}: 1 via FCM, 1 via APNs",
-                zerver_logger.output[1],
+                zerver_logger.output[2],
             )
 
     @responses.activate
     @override_settings(ZILENCER_ENABLED=False)
-    def test_success_self_hosted(self, unused_mock: mock.MagicMock) -> None:
+    def test_success_self_hosted(self) -> None:
         self.add_mock_response()
 
         hamlet = self.example_user("hamlet")
@@ -574,5 +676,5 @@ class RemovePushNotificationTest(E2EEPushNotificationTestCase):
             self.assertEqual(
                 "INFO:zerver.lib.push_notifications:"
                 f"Sent E2EE mobile push notifications for user {hamlet.id}: 1 via FCM, 1 via APNs",
-                zerver_logger.output[1],
+                zerver_logger.output[2],
             )

zerver/tests/test_handle_push_notification.py

@@ -146,10 +146,15 @@ class HandlePushNotificationTest(PushNotificationTestCase):
                 ),
             )
 
+            self.assertIn(
+                "INFO:zerver.lib.push_notifications:"
+                f"Skipping E2EE push notifications for user {self.user_profile.id} because there are no registered devices",
+                pn_logger.output,
+            )
+
     @activate_push_notification_service()
     @responses.activate
-    @mock.patch("zerver.lib.push_notifications.send_push_notifications")
-    def test_end_to_end_failure_due_to_no_plan(self, unused_mock: mock.MagicMock) -> None:
+    def test_end_to_end_failure_due_to_no_plan(self) -> None:
         self.add_mock_response()
 
         self.setup_apns_tokens()
@@ -191,6 +196,7 @@ class HandlePushNotificationTest(PushNotificationTestCase):
                 [
                     f"INFO:zerver.lib.push_notifications:Sending push notifications to mobile clients for user {self.user_profile.id}",
                     "WARNING:zerver.lib.push_notifications:Bouncer refused to send push notification: Your plan doesn't allow sending push notifications. Reason provided by the server: Push notifications access with 10+ users requires signing up for a plan. https://zulip.com/plans/",
+                    f"INFO:zerver.lib.push_notifications:Skipping E2EE push notifications for user {self.user_profile.id} because there are no registered devices",
                 ],
             )
             realm.refresh_from_db()
@@ -215,7 +221,7 @@ class HandlePushNotificationTest(PushNotificationTestCase):
             handle_push_notification(self.user_profile.id, new_missed_message)
             self.assertIn(
                 f"Sent mobile push notifications for user {self.user_profile.id}",
-                pn_logger.output[-1],
+                pn_logger.output[-2],
             )
             realm.refresh_from_db()
             self.assertEqual(realm.push_notifications_enabled, True)
@@ -484,8 +490,7 @@ class HandlePushNotificationTest(PushNotificationTestCase):
                 ],
             )
 
-    @mock.patch("zerver.lib.push_notifications.send_push_notifications")
-    def test_send_notifications_to_bouncer(self, unused_mock: mock.MagicMock) -> None:
+    def test_send_notifications_to_bouncer(self) -> None:
         self.setup_apns_tokens()
         self.setup_fcm_tokens()
 
@@ -554,6 +559,7 @@ class HandlePushNotificationTest(PushNotificationTestCase):
                 [
                     f"INFO:zerver.lib.push_notifications:Sending push notifications to mobile clients for user {user_profile.id}",
                     f"INFO:zerver.lib.push_notifications:Sent mobile push notifications for user {user_profile.id} through bouncer: 3 via FCM devices, 5 via APNs devices",
+                    f"INFO:zerver.lib.push_notifications:Skipping E2EE push notifications for user {self.user_profile.id} because there are no registered devices",
                 ],
             )