zilencer: Add new mobile_pushes_received::day LoggingCountStat.

analytics/lib/counts.py

@@ -23,6 +23,9 @@ from zerver.lib.logging_util import log_to_file
 from zerver.lib.timestamp import ceiling_to_day, ceiling_to_hour, floor_to_hour, verify_UTC
 from zerver.models import Message, Realm, RealmAuditLog, Stream, UserActivityInterval, UserProfile
 
+if settings.ZILENCER_ENABLED:
+    from zilencer.models import RemoteInstallationCount, RemoteZulipServer
+
 ## Logging setup ##
 
 logger = logging.getLogger("zulip.management")
@@ -293,7 +296,7 @@ def do_aggregate_to_summary_table(
 
 # called from zerver.actions; should not throw any errors
 def do_increment_logging_stat(
-    model_object_for_bucket: Union[Realm, UserProfile, Stream],
+    model_object_for_bucket: Union[Realm, UserProfile, Stream, "RemoteZulipServer"],
     stat: CountStat,
     subgroup: Optional[Union[str, int, bool]],
     event_time: datetime,
@@ -305,13 +308,20 @@ def do_increment_logging_stat(
     table = stat.data_collector.output_table
     if table == RealmCount:
         assert isinstance(model_object_for_bucket, Realm)
-        id_args: Dict[str, Union[Realm, UserProfile, Stream]] = {"realm": model_object_for_bucket}
+        id_args: Dict[str, Optional[Union[Realm, UserProfile, Stream, "RemoteZulipServer"]]] = {
+            "realm": model_object_for_bucket
+        }
     elif table == UserCount:
         assert isinstance(model_object_for_bucket, UserProfile)
         id_args = {"realm": model_object_for_bucket.realm, "user": model_object_for_bucket}
-    else:  # StreamCount
+    elif table == StreamCount:
         assert isinstance(model_object_for_bucket, Stream)
         id_args = {"realm": model_object_for_bucket.realm, "stream": model_object_for_bucket}
+    elif table == RemoteInstallationCount:
+        assert isinstance(model_object_for_bucket, RemoteZulipServer)
+        id_args = {"server": model_object_for_bucket, "remote_id": None}
+    else:
+        raise AssertionError("Unsupported CountStat output_table")
 
     if stat.frequency == CountStat.DAY:
         end_time = ceiling_to_day(event_time)
@@ -833,6 +843,15 @@ def get_count_stats(realm: Optional[Realm] = None) -> Dict[str, CountStat]:
         ),
     ]
 
+    if settings.ZILENCER_ENABLED:
+        count_stats_.append(
+            LoggingCountStat(
+                "mobile_pushes_received::day",
+                RemoteInstallationCount,
+                CountStat.DAY,
+            )
+        )
+
     return OrderedDict((stat.property, stat) for stat in count_stats_)
 
 

analytics/tests/test_counts.py

@@ -3,9 +3,11 @@ from typing import Any, Dict, List, Optional, Tuple, Type
 from unittest import mock
 
 import orjson
+import time_machine
 from django.apps import apps
 from django.db import models
 from django.db.models import Sum
+from django.test import override_settings
 from django.utils.timezone import now as timezone_now
 from psycopg2.sql import SQL, Literal
 from typing_extensions import override
@@ -53,14 +55,20 @@ from zerver.actions.user_activity import update_user_activity_interval
 from zerver.actions.users import do_deactivate_user
 from zerver.lib.create_user import create_user
 from zerver.lib.exceptions import InvitationError
+from zerver.lib.push_notifications import (
+    get_message_payload_apns,
+    get_message_payload_gcm,
+    hex_to_b64,
+)
 from zerver.lib.test_classes import ZulipTestCase
-from zerver.lib.timestamp import TimeZoneNotUTCError, floor_to_day
+from zerver.lib.timestamp import TimeZoneNotUTCError, ceiling_to_day, floor_to_day
 from zerver.lib.topic import DB_TOPIC_NAME
 from zerver.lib.utils import assert_is_not_none
 from zerver.models import (
     Client,
     Huddle,
     Message,
+    NotificationTriggers,
     PreregistrationUser,
     Realm,
     RealmAuditLog,
@@ -74,7 +82,7 @@ from zerver.models import (
     get_user,
     is_cross_realm_bot_email,
 )
-from zilencer.models import RemoteInstallationCount, RemoteZulipServer
+from zilencer.models import RemoteInstallationCount, RemotePushDeviceToken, RemoteZulipServer
 from zilencer.views import get_last_id_from_server
 
 
@@ -236,7 +244,7 @@ class AnalyticsTestCase(ZulipTestCase):
                 kwargs[arg_keys[i]] = values[i]
             for key, value in defaults.items():
                 kwargs[key] = kwargs.get(key, value)
-            if table is not InstallationCount and "realm" not in kwargs:
+            if table not in [InstallationCount, RemoteInstallationCount] and "realm" not in kwargs:
                 if "user" in kwargs:
                     kwargs["realm"] = kwargs["user"].realm
                 elif "stream" in kwargs:
@@ -1377,6 +1385,86 @@ class TestLoggingCountStats(AnalyticsTestCase):
             ],
         )
 
+    @override_settings(PUSH_NOTIFICATION_BOUNCER_URL="https://push.zulip.org.example.com")
+    def test_mobile_pushes_received_count(self) -> None:
+        self.server_uuid = "6cde5f7a-1f7e-4978-9716-49f69ebfc9fe"
+        self.server = RemoteZulipServer.objects.create(
+            uuid=self.server_uuid,
+            api_key="magic_secret_api_key",
+            hostname="demo.example.com",
+            last_updated=timezone_now(),
+        )
+
+        hamlet = self.example_user("hamlet")
+        token = "aaaa"
+
+        RemotePushDeviceToken.objects.create(
+            kind=RemotePushDeviceToken.GCM,
+            token=hex_to_b64(token),
+            user_uuid=(hamlet.uuid),
+            server=self.server,
+        )
+        RemotePushDeviceToken.objects.create(
+            kind=RemotePushDeviceToken.GCM,
+            token=hex_to_b64(token + "aa"),
+            user_uuid=(hamlet.uuid),
+            server=self.server,
+        )
+        RemotePushDeviceToken.objects.create(
+            kind=RemotePushDeviceToken.APNS,
+            token=hex_to_b64(token),
+            user_uuid=str(hamlet.uuid),
+            server=self.server,
+        )
+
+        message = Message(
+            sender=hamlet,
+            recipient=self.example_user("othello").recipient,
+            realm_id=hamlet.realm_id,
+            content="This is test content",
+            rendered_content="This is test content",
+            date_sent=timezone_now(),
+            sending_client=get_client("test"),
+        )
+        message.set_topic_name("Test topic")
+        message.save()
+        gcm_payload, gcm_options = get_message_payload_gcm(hamlet, message)
+        apns_payload = get_message_payload_apns(
+            hamlet, message, NotificationTriggers.DIRECT_MESSAGE
+        )
+
+        payload = {
+            "user_id": hamlet.id,
+            "user_uuid": str(hamlet.uuid),
+            "gcm_payload": gcm_payload,
+            "apns_payload": apns_payload,
+            "gcm_options": gcm_options,
+        }
+        now = timezone_now()
+        with time_machine.travel(now, tick=False), mock.patch(
+            "zilencer.views.send_android_push_notification"
+        ), mock.patch("zilencer.views.send_apple_push_notification"), self.assertLogs(
+            "zilencer.views", level="INFO"
+        ):
+            result = self.uuid_post(
+                self.server_uuid,
+                "/api/v1/remotes/push/notify",
+                payload,
+                content_type="application/json",
+                subdomain="",
+            )
+            self.assert_json_success(result)
+
+        # There are 3 devices we created for the user, and the Count increment should
+        # match that number.
+        self.assertTableState(
+            RemoteInstallationCount,
+            ["property", "value", "subgroup", "server", "remote_id", "end_time"],
+            [
+                ["mobile_pushes_received::day", 3, None, self.server, None, ceiling_to_day(now)],
+            ],
+        )
+
     def test_invites_sent(self) -> None:
         property = "invites_sent::day"
 

zilencer/models.py

@@ -135,6 +135,9 @@ class BaseRemoteCount(BaseCount):
     # The remote_id field is the id value of the corresponding *Count object
     # on the remote server.
     # It lets us deduplicate data from the remote server.
+    # Note: Some counts don't come from the remote server, but rather
+    # are stats we track on the bouncer server itself, pertaining to the remote server.
+    # E.g. mobile_pushes_received::day. Such counts will set this field to None.
     remote_id = models.IntegerField(null=True)
 
     class Meta:

zilencer/views.py

@@ -17,7 +17,7 @@ from django.utils.translation import gettext as err_
 from django.views.decorators.csrf import csrf_exempt
 from pydantic import BaseModel, ConfigDict
 
-from analytics.lib.counts import COUNT_STATS
+from analytics.lib.counts import COUNT_STATS, do_increment_logging_stat
 from corporate.lib.stripe import do_deactivate_remote_server
 from zerver.decorator import require_post
 from zerver.lib.exceptions import JsonableError
@@ -397,6 +397,13 @@ def remote_server_notify_push(
         len(android_devices),
         len(apple_devices),
     )
+    do_increment_logging_stat(
+        server,
+        COUNT_STATS["mobile_pushes_received::day"],
+        None,
+        timezone_now(),
+        increment=len(android_devices) + len(apple_devices),
+    )
 
     # Truncate incoming pushes to 200, due to APNs maximum message
     # sizes; see handle_remove_push_notification for the version of