From 82ee93d967fa72894ac526527369613bf7a2fb33 Mon Sep 17 00:00:00 2001
From: Mateusz Mandera <mateusz.mandera@zulip.com>
Date: Fri, 22 Dec 2023 02:22:48 +0100
Subject: [PATCH] zilencer: Set .remote_realm for existing
 RemotePushDeviceToken.

Old RemotePushDeviceTokens were created without this attribute. But when
processing a notification, if we have remote_realm, we can take the
opportunity to to set this for all the registrations for this user.
---
 zerver/tests/test_push_notifications.py | 48 +++++++++++++++++++++++++
 zilencer/views.py                       | 17 +++++++++
 2 files changed, 65 insertions(+)

diff --git a/zerver/tests/test_push_notifications.py b/zerver/tests/test_push_notifications.py
index ec74e192f0..759155852a 100644
--- a/zerver/tests/test_push_notifications.py
+++ b/zerver/tests/test_push_notifications.py
@@ -555,6 +555,54 @@ class PushBouncerNotificationTest(BouncerTestCase):
             result = self.uuid_post(self.server_uuid, endpoint, payload)
             self.assert_json_error(result, "Empty or invalid length token")
 
+    def test_send_notification_endpoint_sets_remote_realm_for_devices(self) -> None:
+        hamlet = self.example_user("hamlet")
+        server = self.server
+
+        remote_realm = RemoteRealm.objects.get(server=server, uuid=hamlet.realm.uuid)
+
+        android_token = RemotePushDeviceToken.objects.create(
+            kind=RemotePushDeviceToken.GCM,
+            token=hex_to_b64("aaaa"),
+            user_uuid=hamlet.uuid,
+            server=server,
+        )
+        apple_token = RemotePushDeviceToken.objects.create(
+            kind=RemotePushDeviceToken.APNS,
+            token=hex_to_b64("bbbb"),
+            user_uuid=hamlet.uuid,
+            server=server,
+        )
+        payload = {
+            "user_id": hamlet.id,
+            "user_uuid": str(hamlet.uuid),
+            "realm_uuid": str(hamlet.realm.uuid),
+            "gcm_payload": {},
+            "apns_payload": {},
+            "gcm_options": {},
+        }
+        with mock.patch(
+            "zilencer.views.send_android_push_notification", return_value=1
+        ), mock.patch("zilencer.views.send_apple_push_notification", return_value=1), mock.patch(
+            "corporate.lib.stripe.RemoteServerBillingSession.current_count_for_billed_licenses",
+            return_value=10,
+        ), self.assertLogs(
+            "zilencer.views", level="INFO"
+        ):
+            result = self.uuid_post(
+                self.server_uuid,
+                "/api/v1/remotes/push/notify",
+                payload,
+                content_type="application/json",
+            )
+        self.assert_json_success(result)
+
+        android_token.refresh_from_db()
+        apple_token.refresh_from_db()
+
+        self.assertEqual(android_token.remote_realm, remote_realm)
+        self.assertEqual(apple_token.remote_realm, remote_realm)
+
     def test_send_notification_endpoint(self) -> None:
         hamlet = self.example_user("hamlet")
         server = RemoteZulipServer.objects.get(uuid=self.server_uuid)
diff --git a/zilencer/views.py b/zilencer/views.py
index 77abaf37e0..28832e857b 100644
--- a/zilencer/views.py
+++ b/zilencer/views.py
@@ -494,6 +494,9 @@ def remote_server_notify_push(
         increment=len(android_devices) + len(apple_devices),
     )
     if remote_realm is not None:
+        ensure_devices_set_remote_realm(
+            android_devices=android_devices, apple_devices=apple_devices, remote_realm=remote_realm
+        )
         do_increment_logging_stat(
             remote_realm,
             COUNT_STATS["mobile_pushes_received::day"],
@@ -671,6 +674,20 @@ def batch_create_table_data(
         )
 
 
+def ensure_devices_set_remote_realm(
+    android_devices: List[RemotePushDeviceToken],
+    apple_devices: List[RemotePushDeviceToken],
+    remote_realm: RemoteRealm,
+) -> None:
+    devices_to_update = []
+    for device in android_devices + apple_devices:
+        if device.remote_realm_id is None:
+            device.remote_realm = remote_realm
+            devices_to_update.append(device)
+
+    RemotePushDeviceToken.objects.bulk_update(devices_to_update, ["remote_realm"])
+
+
 def update_remote_realm_data_for_server(
     server: RemoteZulipServer, server_realms_info: List[RealmDataForAnalytics]
 ) -> None: