diff --git a/corporate/tests/test_remote_billing.py b/corporate/tests/test_remote_billing.py index 3b11bf91e3..8e232eb77e 100644 --- a/corporate/tests/test_remote_billing.py +++ b/corporate/tests/test_remote_billing.py @@ -27,10 +27,12 @@ from zerver.actions.realm_settings import do_deactivate_realm from zerver.lib.exceptions import RemoteRealmServerMismatchError from zerver.lib.rate_limiter import RateLimitedIPAddr from zerver.lib.remote_server import send_server_data_to_push_bouncer +from zerver.lib.send_email import FromAddress from zerver.lib.test_classes import BouncerTestCase from zerver.lib.test_helpers import ratelimit_rule from zerver.lib.timestamp import datetime_to_timestamp from zerver.models import Realm, UserProfile +from zerver.models.realms import get_realm from zilencer.models import ( PreregistrationRemoteRealmBillingUser, PreregistrationRemoteServerBillingUser, @@ -876,6 +878,117 @@ class RemoteBillingAuthenticationTest(RemoteRealmBillingTestCase): RemoteRealmBillingSession(remote_realm_with_plan).get_next_plan(plan), server_next_plan ) + @responses.activate + def test_transfer_plan_from_server_to_realm_when_realm_has_customer( + self, + ) -> None: + self.login("desdemona") + desdemona = self.example_user("desdemona") + zulip_realm = get_realm("zulip") + + server_billing_session = RemoteServerBillingSession(self.server) + server_customer = server_billing_session.update_or_create_customer(stripe_customer_id=None) + server_plan = CustomerPlan.objects.create( + customer=server_customer, + billing_cycle_anchor=timezone_now(), + billing_schedule=CustomerPlan.BILLING_SCHEDULE_ANNUAL, + tier=CustomerPlan.TIER_SELF_HOSTED_COMMUNITY, + status=CustomerPlan.ACTIVE, + ) + self.server.plan_type = RemoteZulipServer.PLAN_TYPE_COMMUNITY + self.server.save(update_fields=["plan_type"]) + + # Delete any existing remote realms. + RemoteRealm.objects.all().delete() + + # We want there to be only a single (non-system bot) realm on the server for our setup. + Realm.objects.exclude(string_id__in=["zulip", "zulipinternal"]).update(deactivated=True) + + # Send server data to push bouncer. + self.add_mock_response() + send_server_data_to_push_bouncer(consider_usage_statistics=False) + + # Let's create a plan for the realm. This will conflict with the server plan. + remote_realm = RemoteRealm.objects.get(uuid=zulip_realm.uuid) + realm_billing_session = RemoteRealmBillingSession(remote_realm) + realm_customer = realm_billing_session.update_or_create_customer(stripe_customer_id=None) + realm_plan = CustomerPlan.objects.create( + customer=realm_customer, + billing_cycle_anchor=timezone_now(), + billing_schedule=CustomerPlan.BILLING_SCHEDULE_ANNUAL, + tier=CustomerPlan.TIER_SELF_HOSTED_LEGACY, + status=CustomerPlan.ACTIVE, + ) + remote_realm.plan_type = RemoteRealm.PLAN_TYPE_SELF_MANAGED_LEGACY + remote_realm.save(update_fields=["plan_type"]) + + with self.assertLogs("zilencer.views", "WARN") as mock_warn: + result = self.execute_remote_billing_authentication_flow( + desdemona, return_from_auth_url=True + ) + self.assertEqual( + mock_warn.output, + [ + f"WARNING:zilencer.views:Failed to migrate customer from server (id: {remote_realm.server.id}) to realm (id: {remote_realm.id}): " + "RemoteRealm customer already exists and plans can't be migrated automatically." + ], + ) + self.assert_json_error( + result, + f"Couldn't reconcile billing data between server and realm. Please contact {FromAddress.SUPPORT}", + ) + + # If the realm's plan is ENDED, it's safe to move the server plan over. + realm_plan.status = CustomerPlan.ENDED + realm_plan.save(update_fields=["status"]) + # However, not if the server's status indicates that there's some kind + # of plan change queued up after the plan, since that state would be + # harder and more risky to try to migrate. + server_plan.status = CustomerPlan.SWITCH_PLAN_TIER_AT_PLAN_END + server_plan.save(update_fields=["status"]) + + with self.assertLogs("zilencer.views", "WARN") as mock_warn: + result = self.execute_remote_billing_authentication_flow( + desdemona, return_from_auth_url=True + ) + self.assertEqual( + mock_warn.output, + [ + f"WARNING:zilencer.views:Failed to migrate customer from server (id: {remote_realm.server.id}) to realm (id: {remote_realm.id}): " + "RemoteRealm customer already exists and plans can't be migrated automatically." + ], + ) + self.assert_json_error( + result, + f"Couldn't reconcile billing data between server and realm. Please contact {FromAddress.SUPPORT}", + ) + + # Finally, we simulate a regular, ACTIVE plan for the server again. Combined with + # the ENDED plan for the realm, we now have a simple case, where the migration + # should proceed. + server_plan.status = CustomerPlan.ACTIVE + server_plan.save(update_fields=["status"]) + + result = self.execute_remote_billing_authentication_flow( + desdemona, return_from_auth_url=False + ) + self.assertEqual(result.status_code, 302) + + # Server plan status was reset + self.server.refresh_from_db() + self.assertEqual(self.server.plan_type, RemoteZulipServer.PLAN_TYPE_SELF_MANAGED) + + # The Customer objects remain as they were. + self.assertEqual(get_customer_by_remote_realm(remote_realm), realm_customer) + self.assertEqual(get_customer_by_remote_server(self.server), server_customer) + + # The plan that used to be for the server, has been migrated to the realm customer: + self.assertEqual(get_current_plan_by_customer(server_customer), None) + self.assertEqual(get_current_plan_by_customer(realm_customer), server_plan) + + remote_realm.refresh_from_db() + self.assertEqual(remote_realm.plan_type, RemoteRealm.PLAN_TYPE_COMMUNITY) + @responses.activate def test_transfer_business_plan_from_server_to_realm( self, diff --git a/corporate/views/remote_billing_page.py b/corporate/views/remote_billing_page.py index 2890713585..7ca7ffe25a 100644 --- a/corporate/views/remote_billing_page.py +++ b/corporate/views/remote_billing_page.py @@ -201,6 +201,10 @@ def remote_realm_billing_finalize_login( try: handle_customer_migration_from_server_to_realms(server=remote_server) + except JsonableError: + # JsonableError should be propagated up, as they are meant to convey + # a json error response to be returned. + raise except Exception: # nocoverage billing_logger.exception( "%s: Failed to migrate customer from server (id: %s) to realms", diff --git a/zilencer/views.py b/zilencer/views.py index b6d051cf2c..23eb6ca30a 100644 --- a/zilencer/views.py +++ b/zilencer/views.py @@ -32,7 +32,11 @@ from corporate.lib.stripe import ( do_deactivate_remote_server, get_push_status_for_remote_request, ) -from corporate.models import CustomerPlan, get_current_plan_by_customer +from corporate.models import ( + CustomerPlan, + get_current_plan_by_customer, + get_customer_by_remote_realm, +) from zerver.decorator import require_post from zerver.lib.email_validation import validate_disposable from zerver.lib.exceptions import ( @@ -56,6 +60,7 @@ from zerver.lib.remote_server import ( ) from zerver.lib.request import REQ, RequestNotes, has_request_variables from zerver.lib.response import json_success +from zerver.lib.send_email import FromAddress from zerver.lib.timestamp import timestamp_to_datetime from zerver.lib.typed_endpoint import JsonBodyPayload, typed_endpoint from zerver.lib.types import RemoteRealmDictValue @@ -1000,13 +1005,47 @@ def handle_customer_migration_from_server_to_realms( elif len(realm_uuids) == 1: # Here, we have exactly one non-system-bot realm, and some # sort of plan on the server; move it to the realm. - remote_realm = RemoteRealm.objects.get( - uuid=realm_uuids[0], plan_type=RemoteRealm.PLAN_TYPE_SELF_MANAGED - ) + remote_realm = RemoteRealm.objects.get(uuid=realm_uuids[0], server=server) + remote_realm_customer = get_customer_by_remote_realm(remote_realm) + # Migrate customer from server to remote realm if there is only one realm. - server_customer.remote_realm = remote_realm - server_customer.remote_server = None - server_customer.save(update_fields=["remote_realm", "remote_server"]) + if remote_realm_customer is None: + # In this case the migration is easy, since we can just move the customer + # object directly. + server_customer.remote_realm = remote_realm + server_customer.remote_server = None + server_customer.save(update_fields=["remote_realm", "remote_server"]) + else: + # If there's a Customer object for the realm already, things are harder, + # because it's an unusual state and there may be a plan already active + # for the realm, or there may have been. + # In the simplest case, where the realm doesn't have an active plan and the + # server's plan state can easily be moved, we proceed with the migrations. + remote_realm_plan = get_current_plan_by_customer(remote_realm_customer) + if ( + remote_realm_plan is None + and server_plan.status != CustomerPlan.SWITCH_PLAN_TIER_AT_PLAN_END + ): + # This is a simple case where we don't have to worry about the realm already + # having an active plan, or the server having a next plan scheduled that we'd need + # to figure out how to migrate correctly as well. + # Any other case is too complex to handle here, and should be handled manually, + # especially since that should be extremely rare. + server_plan.customer = remote_realm_customer + server_plan.save(update_fields=["customer"]) + else: + logger.warning( + "Failed to migrate customer from server (id: %s) to realm (id: %s): RemoteRealm customer already exists " + "and plans can't be migrated automatically.", + server.id, + remote_realm.id, + ) + raise JsonableError( + _( + "Couldn't reconcile billing data between server and realm. Please contact {support_email}" + ).format(support_email=FromAddress.SUPPORT) + ) + # TODO: Might be better to call do_change_plan_type here. remote_realm.plan_type = server.plan_type remote_realm.save(update_fields=["plan_type"])