remote_server_post_analytics: Return remote realms data in response.

This is a prep commit to return, for each remote realm, the 'uuid',
'can_push', and 'expected_end_timestamp'.

This data will be used in 'initialize_push_notifications'.

zerver/lib/remote_server.py

@@ -15,6 +15,7 @@ from zerver.lib.exceptions import JsonableError, MissingRemoteRealmError
 from zerver.lib.export import floatify_datetime_fields
 from zerver.lib.outgoing_http import OutgoingSession
 from zerver.lib.queue import queue_json_publish
+from zerver.lib.types import RemoteRealmDictValue
 from zerver.models import OrgTypeEnum, Realm, RealmAuditLog
 
 
@@ -287,7 +288,7 @@ def send_analytics_to_push_bouncer() -> None:
     logger.info("Reported %d records", record_count)
 
 
-def send_realms_only_to_push_bouncer() -> None:
+def send_realms_only_to_push_bouncer() -> Dict[str, RemoteRealmDictValue]:
     request = {
         "realm_counts": "[]",
         "installation_counts": "[]",
@@ -299,7 +300,10 @@ def send_realms_only_to_push_bouncer() -> None:
 
     # 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)
+    response = send_to_push_bouncer("POST", "server/analytics", request)
+    assert isinstance(response["realms"], dict)  # for mypy
+
+    return response["realms"]
 
 
 def enqueue_register_realm_with_push_bouncer_if_needed(realm: Realm) -> None:

zerver/lib/types.py

@@ -315,3 +315,8 @@ class RawUserDict(TypedDict):
     bot_type: Optional[int]
     long_term_idle: bool
     email_address_visibility: int
+
+
+class RemoteRealmDictValue(TypedDict):
+    can_push: bool
+    expected_end_timestamp: Optional[int]

zerver/tests/test_push_notifications.py

@@ -23,6 +23,7 @@ from typing_extensions import override
 
 from analytics.lib.counts import CountStat, LoggingCountStat
 from analytics.models import InstallationCount, RealmCount
+from corporate.models import CustomerPlan
 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
@@ -1577,6 +1578,47 @@ class AnalyticsBouncerTest(BouncerTestCase):
     def test_send_realms_only_to_push_bouncer(self) -> None:
         self.add_mock_response()
 
+        with mock.patch(
+            "zilencer.views.RemoteRealmBillingSession.get_customer", return_value=None
+        ) as m:
+            realms = send_realms_only_to_push_bouncer()
+            m.assert_called()
+            for data in realms.values():
+                self.assertEqual(data["can_push"], True)
+                self.assertEqual(data["expected_end_timestamp"], None)
+
+        dummy_customer = mock.MagicMock()
+        with mock.patch(
+            "zilencer.views.RemoteRealmBillingSession.get_customer", return_value=dummy_customer
+        ):
+            with mock.patch("zilencer.views.get_current_plan_by_customer", return_value=None) as m:
+                realms = send_realms_only_to_push_bouncer()
+                m.assert_called()
+                for data in realms.values():
+                    self.assertEqual(data["can_push"], True)
+                    self.assertEqual(data["expected_end_timestamp"], None)
+
+        dummy_customer_plan = mock.MagicMock()
+        dummy_customer_plan.status = CustomerPlan.DOWNGRADE_AT_END_OF_CYCLE
+        dummy_date = datetime.datetime(year=2023, month=12, day=3, tzinfo=datetime.timezone.utc)
+        with mock.patch(
+            "zilencer.views.RemoteRealmBillingSession.get_customer", return_value=dummy_customer
+        ):
+            with mock.patch(
+                "zilencer.views.get_current_plan_by_customer", return_value=dummy_customer_plan
+            ):
+                with mock.patch(
+                    "zilencer.views.RemoteRealmBillingSession.get_next_billing_cycle",
+                    return_value=dummy_date,
+                ) as m:
+                    realms = send_realms_only_to_push_bouncer()
+                    m.assert_called()
+                    for data in realms.values():
+                        self.assertEqual(data["can_push"], True)
+                        self.assertEqual(
+                            data["expected_end_timestamp"], datetime_to_timestamp(dummy_date)
+                        )
+
         send_realms_only_to_push_bouncer()
 
         self.assertEqual(
@@ -1608,8 +1650,19 @@ class AnalyticsBouncerTest(BouncerTestCase):
         )
 
         # Use a mock to assert exactly the data that gets sent.
+        dummy_send_realms_only_response = {
+            "result": "success",
+            "msg": "",
+            "realms": {
+                "f9535515-84d0-489e-80d5-9ae97c3c7ec1": {
+                    "can_push": True,
+                    "expected_end_timestamp": None,
+                },
+            },
+        }
         with mock.patch(
-            "zerver.lib.remote_server.send_to_push_bouncer"
+            "zerver.lib.remote_server.send_to_push_bouncer",
+            return_value=dummy_send_realms_only_response,
         ) as mock_send_to_push_bouncer:
             send_realms_only_to_push_bouncer()
 

zerver/tests/test_realm.py

@@ -1065,7 +1065,20 @@ class RealmTest(ZulipTestCase):
 
     @override_settings(PUSH_NOTIFICATION_BOUNCER_URL="https://push.zulip.org.example.com")
     def test_do_create_realm_notify_bouncer(self) -> None:
-        with mock.patch("zerver.lib.remote_server.send_to_push_bouncer") as m:
+        dummy_send_realms_only_response = {
+            "result": "success",
+            "msg": "",
+            "realms": {
+                "dummy-uuid": {
+                    "can_push": True,
+                    "expected_end_timestamp": None,
+                },
+            },
+        }
+        with mock.patch(
+            "zerver.lib.remote_server.send_to_push_bouncer",
+            return_value=dummy_send_realms_only_response,
+        ) as m:
             realm = do_create_realm("realm_string_id", "realm name")
 
         self.assertEqual(realm.string_id, "realm_string_id")

zilencer/views.py

@@ -23,7 +23,8 @@ from analytics.lib.counts import (
     REMOTE_INSTALLATION_COUNT_STATS,
     do_increment_logging_stat,
 )
-from corporate.lib.stripe import do_deactivate_remote_server
+from corporate.lib.stripe import RemoteRealmBillingSession, do_deactivate_remote_server
+from corporate.models import CustomerPlan, get_current_plan_by_customer
 from zerver.decorator import require_post
 from zerver.lib.exceptions import JsonableError
 from zerver.lib.push_notifications import (
@@ -36,8 +37,9 @@ from zerver.lib.push_notifications import (
 from zerver.lib.remote_server import RealmDataForAnalytics
 from zerver.lib.request import REQ, has_request_variables
 from zerver.lib.response import json_success
-from zerver.lib.timestamp import timestamp_to_datetime
+from zerver.lib.timestamp import datetime_to_timestamp, timestamp_to_datetime
 from zerver.lib.typed_endpoint import JsonBodyPayload, typed_endpoint
+from zerver.lib.types import RemoteRealmDictValue
 from zerver.lib.validator import check_capped_string, check_int, check_string_fixed_length
 from zerver.views.push_notifications import check_app_id, validate_token
 from zilencer.auth import InvalidZulipServerKeyError
@@ -766,7 +768,36 @@ def remote_server_post_analytics(
             )
         batch_create_table_data(server, RemoteRealmAuditLog, remote_realm_audit_logs)
 
-    return json_success(request)
+    remote_realm_dict: Dict[str, RemoteRealmDictValue] = {}
+    remote_realms = RemoteRealm.objects.filter(server=server)
+    for remote_realm in remote_realms:
+        uuid = str(remote_realm.uuid)
+        billing_session = RemoteRealmBillingSession(remote_realm)
+
+        customer = billing_session.get_customer()
+        if customer is None:
+            remote_realm_dict[uuid] = {"can_push": True, "expected_end_timestamp": None}
+            continue
+
+        current_plan = get_current_plan_by_customer(customer)
+        if current_plan is None:
+            remote_realm_dict[uuid] = {"can_push": True, "expected_end_timestamp": None}
+            continue
+
+        expected_end_timestamp = None
+        if current_plan.status in [
+            CustomerPlan.DOWNGRADE_AT_END_OF_CYCLE,
+            CustomerPlan.DOWNGRADE_AT_END_OF_FREE_TRIAL,
+        ]:
+            expected_end_timestamp = datetime_to_timestamp(
+                billing_session.get_next_billing_cycle(current_plan)
+            )
+        remote_realm_dict[uuid] = {
+            "can_push": True,
+            "expected_end_timestamp": expected_end_timestamp,
+        }
+
+    return json_success(request, data={"realms": remote_realm_dict})
 
 
 def get_last_id_from_server(server: RemoteZulipServer, model: Any) -> int: