diff --git a/zerver/lib/remote_server.py b/zerver/lib/remote_server.py index 245d8cc663..dace0e259c 100644 --- a/zerver/lib/remote_server.py +++ b/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: diff --git a/zerver/lib/types.py b/zerver/lib/types.py index 764c9d4890..7d79a5b701 100644 --- a/zerver/lib/types.py +++ b/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] diff --git a/zerver/tests/test_push_notifications.py b/zerver/tests/test_push_notifications.py index 0fa86c06e1..a958fd3a94 100644 --- a/zerver/tests/test_push_notifications.py +++ b/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() diff --git a/zerver/tests/test_realm.py b/zerver/tests/test_realm.py index 80b242023e..094b2a3ac1 100644 --- a/zerver/tests/test_realm.py +++ b/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") diff --git a/zilencer/views.py b/zilencer/views.py index a8c4bf1ba7..2bef6f10c3 100644 --- a/zilencer/views.py +++ b/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: