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'.
This commit is contained in:
Prakhar Pratyush 2023-11-29 21:30:19 +05:30 committed by Tim Abbott
parent 895d76f6f0
commit 6aa911a9b2
5 changed files with 113 additions and 7 deletions

View File

@ -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:

View File

@ -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]

View File

@ -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()

View File

@ -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")

View File

@ -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: