From 5dc1d36f730d6ce5d5aa47d9bfbf1da0341cef28 Mon Sep 17 00:00:00 2001
From: Tim Abbott <tabbott@zulip.com>
Date: Sat, 16 Dec 2023 21:59:16 -0800
Subject: [PATCH] zilencer: Fix community plan push notifications.

---
 corporate/lib/stripe.py                 |  7 +++
 zerver/tests/test_push_notifications.py | 61 ++++++++++++++++++++++++-
 2 files changed, 67 insertions(+), 1 deletion(-)

diff --git a/corporate/lib/stripe.py b/corporate/lib/stripe.py
index e2e08a2d35..221d2e3ba3 100644
--- a/corporate/lib/stripe.py
+++ b/corporate/lib/stripe.py
@@ -4091,6 +4091,13 @@ def get_push_status_for_remote_request(
     else:
         current_plan = None
 
+    if billing_session.is_sponsored():
+        return PushNotificationsEnabledStatus(
+            can_push=True,
+            expected_end_timestamp=None,
+            message="Community plan",
+        )
+
     if current_plan is not None:
         if current_plan.status in [
             CustomerPlan.DOWNGRADE_AT_END_OF_CYCLE,
diff --git a/zerver/tests/test_push_notifications.py b/zerver/tests/test_push_notifications.py
index 5a624e3e6b..397fd7cb00 100644
--- a/zerver/tests/test_push_notifications.py
+++ b/zerver/tests/test_push_notifications.py
@@ -663,7 +663,7 @@ class PushBouncerNotificationTest(BouncerTestCase):
             remote=server,
         )
 
-    def test_send_notification_endpoint_on_self_managed_plan(self) -> None:
+    def test_send_notification_endpoint_on_free_plans(self) -> None:
         hamlet = self.example_user("hamlet")
         remote_server = RemoteZulipServer.objects.get(uuid=self.server_uuid)
         RemotePushDeviceToken.objects.create(
@@ -791,6 +791,39 @@ class PushBouncerNotificationTest(BouncerTestCase):
         self.assert_json_error(result, "Your plan doesn't allow sending push notifications.")
         self.assertEqual(orjson.loads(result.content)["code"], "BAD_REQUEST")
 
+        # Check that sponsored realms are allowed to send push notifications.
+        remote_server.plan_type = RemoteRealm.PLAN_TYPE_COMMUNITY
+        remote_server.save()
+        with self.assertLogs("zilencer.views", level="INFO") as logger:
+            result = self.uuid_post(
+                self.server_uuid,
+                "/api/v1/remotes/push/notify",
+                payload,
+                content_type="application/json",
+            )
+        data = self.assert_json_success(result)
+        self.assertEqual(
+            {
+                "result": "success",
+                "msg": "",
+                "realm": None,
+                "total_android_devices": 1,
+                "total_apple_devices": 0,
+                "deleted_devices": {"android_devices": [], "apple_devices": []},
+            },
+            data,
+        )
+        self.assertIn(
+            "INFO:zilencer.views:"
+            f"Sending mobile push notifications for remote user 6cde5f7a-1f7e-4978-9716-49f69ebfc9fe:<id:{hamlet.id}><uuid:{hamlet.uuid}>: "
+            "1 via FCM devices, 0 via APNs devices",
+            logger.output,
+        )
+
+        # Reset the plan_type to test remaining cases.
+        remote_server.plan_type = RemoteRealm.PLAN_TYPE_SELF_MANAGED
+        remote_server.save()
+
         human_counts = {
             str(UserProfile.ROLE_REALM_ADMINISTRATOR): 1,
             str(UserProfile.ROLE_REALM_OWNER): 1,
@@ -2152,6 +2185,32 @@ class AnalyticsBouncerTest(BouncerTestCase):
                         self.assertEqual(realm.push_notifications_enabled, False)
                         self.assertEqual(realm.push_notifications_enabled_end_timestamp, None)
 
+        RemoteRealm.objects.filter(server=self.server).update(
+            plan_type=RemoteRealm.PLAN_TYPE_COMMUNITY
+        )
+
+        with mock.patch(
+            "zilencer.views.RemoteRealmBillingSession.get_customer", return_value=dummy_customer
+        ) as m:
+            with mock.patch(
+                "corporate.lib.stripe.get_current_plan_by_customer", return_value=None
+            ) as m:
+                with mock.patch(
+                    "corporate.lib.stripe.RemoteRealmBillingSession.current_count_for_billed_licenses",
+                    return_value=11,
+                ):
+                    send_server_data_to_push_bouncer(consider_usage_statistics=False)
+                    m.assert_called()
+                    realms = Realm.objects.all()
+                    for realm in realms:
+                        self.assertEqual(realm.push_notifications_enabled, True)
+                        self.assertEqual(realm.push_notifications_enabled_end_timestamp, None)
+
+        # Reset the plan type to test remaining cases.
+        RemoteRealm.objects.filter(server=self.server).update(
+            plan_type=RemoteRealm.PLAN_TYPE_SELF_MANAGED
+        )
+
         dummy_customer_plan = mock.MagicMock()
         dummy_customer_plan.status = CustomerPlan.DOWNGRADE_AT_END_OF_CYCLE
         dummy_date = datetime(year=2023, month=12, day=3, tzinfo=timezone.utc)