From ab633f45575fc115232c04e9f623dcc8caa375d7 Mon Sep 17 00:00:00 2001
From: Mateusz Mandera <mateusz.mandera@zulip.com>
Date: Thu, 16 Nov 2023 15:25:58 +0100
Subject: [PATCH] analytics: Add send_realms_only_to_push_bouncer function.

This is a useful helper using the same API as
send_analytics_to_push_bouncer(), but uploading only realms info. This
is useful to upload realms info without the risk of taking a long time
to process the request due to too much of the *Count analytics data.
---
 zerver/lib/remote_server.py             | 13 ++++++
 zerver/tests/test_push_notifications.py | 57 +++++++++++++++++++++++++
 2 files changed, 70 insertions(+)

diff --git a/zerver/lib/remote_server.py b/zerver/lib/remote_server.py
index 5aad3419b3..89b5bb771a 100644
--- a/zerver/lib/remote_server.py
+++ b/zerver/lib/remote_server.py
@@ -229,3 +229,16 @@ def send_analytics_to_push_bouncer() -> None:
         send_to_push_bouncer("POST", "server/analytics", request)
     except JsonableError as e:
         logging.warning(e.msg)
+
+
+def send_realms_only_to_push_bouncer() -> None:
+    request = {
+        "realm_counts": "[]",
+        "installation_counts": "[]",
+        "realms": orjson.dumps(get_realms_info_for_push_bouncer()).decode(),
+        "version": orjson.dumps(ZULIP_VERSION).decode(),
+    }
+
+    # We don't catch JsonableError here, because we want it to propagate further
+    # to either explicitly, loudly fail or be error-handled by the caller.
+    send_to_push_bouncer("POST", "server/analytics", request)
diff --git a/zerver/tests/test_push_notifications.py b/zerver/tests/test_push_notifications.py
index d46b676d5e..66fb08bac6 100644
--- a/zerver/tests/test_push_notifications.py
+++ b/zerver/tests/test_push_notifications.py
@@ -25,6 +25,7 @@ from typing_extensions import override
 
 from analytics.lib.counts import CountStat, LoggingCountStat
 from analytics.models import InstallationCount, RealmCount
+from version import ZULIP_VERSION
 from zerver.actions.message_delete import do_delete_messages
 from zerver.actions.message_flags import do_mark_stream_messages_as_read, do_update_message_flags
 from zerver.actions.realm_settings import do_deactivate_realm
@@ -59,7 +60,9 @@ from zerver.lib.remote_server import (
     PushNotificationBouncerError,
     PushNotificationBouncerRetryLaterError,
     build_analytics_data,
+    get_realms_info_for_push_bouncer,
     send_analytics_to_push_bouncer,
+    send_realms_only_to_push_bouncer,
     send_to_push_bouncer,
 )
 from zerver.lib.response import json_response_from_error
@@ -1445,6 +1448,60 @@ class AnalyticsBouncerTest(BouncerTestCase):
             )
         self.assertIn("Malformed audit log data", m.output[0])
 
+    @override_settings(PUSH_NOTIFICATION_BOUNCER_URL="https://push.zulip.org.example.com")
+    @responses.activate
+    def test_send_realms_only_to_push_bouncer(self) -> None:
+        self.add_mock_response()
+        self.example_user("hamlet")
+
+        send_realms_only_to_push_bouncer()
+
+        self.assertEqual(
+            list(
+                RemoteRealm.objects.order_by("id").values(
+                    "server_id",
+                    "uuid",
+                    "uuid_owner_secret",
+                    "host",
+                    "realm_date_created",
+                    "registration_deactivated",
+                    "realm_deactivated",
+                    "plan_type",
+                )
+            ),
+            [
+                {
+                    "server_id": self.server.id,
+                    "uuid": realm.uuid,
+                    "uuid_owner_secret": realm.uuid_owner_secret,
+                    "host": realm.host,
+                    "realm_date_created": realm.date_created,
+                    "registration_deactivated": False,
+                    "realm_deactivated": False,
+                    "plan_type": RemoteRealm.PLAN_TYPE_SELF_HOSTED,
+                }
+                for realm in Realm.objects.order_by("id")
+            ],
+        )
+
+        # Use a mock to assert exactly the data that gets sent.
+        with mock.patch(
+            "zerver.lib.remote_server.send_to_push_bouncer"
+        ) as mock_send_to_push_bouncer:
+            send_realms_only_to_push_bouncer()
+
+        post_data = {
+            "realm_counts": "[]",
+            "installation_counts": "[]",
+            "realms": orjson.dumps(get_realms_info_for_push_bouncer()).decode(),
+            "version": orjson.dumps(ZULIP_VERSION).decode(),
+        }
+        mock_send_to_push_bouncer.assert_called_with(
+            "POST",
+            "server/analytics",
+            post_data,
+        )
+
 
 class PushNotificationTest(BouncerTestCase):
     @override