From 4715a058b0ace97e92f360ac1cbf169648e3441d Mon Sep 17 00:00:00 2001 From: Prakhar Pratyush Date: Fri, 5 Jan 2024 18:54:15 +0530 Subject: [PATCH] migrate_customers: Migrate customer from server to realms during login. Earlier, the 'handle_customer_migration_from_server_to_realms' function was called during the send analytics step. It resulted in an error for customers having multiple Zulip servers, one for testing and the others for not-testing, sharing a push bouncer registration. The migration step when run in a test instance caused customers to have their legacy plan migrated to a test realm, resulting in them losing their legacy plan. This commit moves the migration step to run during plan management login step. This reduces the chances of losing legacy plan as we expect them to only verify that 8.0 upgrade works and not bother trying to login to plan management from their test instance. --- corporate/tests/test_remote_billing.py | 77 ++++++++++++++++++++++++-- corporate/tests/test_stripe.py | 20 +++---- corporate/views/remote_billing_page.py | 24 +++++++- zilencer/views.py | 15 ----- 4 files changed, 106 insertions(+), 30 deletions(-) diff --git a/corporate/tests/test_remote_billing.py b/corporate/tests/test_remote_billing.py index 097fc3ab34..1a199255e6 100644 --- a/corporate/tests/test_remote_billing.py +++ b/corporate/tests/test_remote_billing.py @@ -51,6 +51,7 @@ class RemoteRealmBillingTestCase(BouncerTestCase): # This only matters if first_time_login is True, since otherwise # there's no confirmation link to be clicked: return_without_clicking_confirmation_link: bool = False, + server_on_active_plan_error: bool = False, ) -> "TestHttpResponse": now = timezone_now() @@ -68,6 +69,10 @@ class RemoteRealmBillingTestCase(BouncerTestCase): with time_machine.travel(now, tick=False): result = self.client_get(signed_auth_url, subdomain="selfhosting") + if server_on_active_plan_error: + self.assert_in_response("Plan management not available", result) + return result + if first_time_login: self.assertFalse(RemoteRealmBillingUser.objects.filter(user_uuid=user.uuid).exists()) # When logging in for the first time some extra steps are needed @@ -463,6 +468,9 @@ class RemoteBillingAuthenticationTest(RemoteRealmBillingTestCase): @responses.activate def test_transfer_legacy_plan_from_server_to_all_realms(self) -> None: + self.login("desdemona") + desdemona = self.example_user("desdemona") + # Assert current server is not on any plan. self.assertIsNone(get_customer_by_remote_server(self.server)) @@ -497,6 +505,12 @@ class RemoteBillingAuthenticationTest(RemoteRealmBillingTestCase): self.add_mock_response() send_server_data_to_push_bouncer(consider_usage_statistics=False) + # Login to plan management. + result = self.execute_remote_billing_authentication_flow( + desdemona, server_on_active_plan_error=True + ) + self.assertEqual(result.status_code, 200) + # RemoteRealm objects should be created for all realms on the server. self.assert_length(RemoteRealm.objects.all(), 4) @@ -509,8 +523,11 @@ class RemoteBillingAuthenticationTest(RemoteRealmBillingTestCase): server_customer.sponsorship_pending = False server_customer.save() - # Send server data to push bouncer again. - send_server_data_to_push_bouncer(consider_usage_statistics=False) + # Login to plan management. Performs customer migration from server to realms. + result = self.execute_remote_billing_authentication_flow( + desdemona, server_on_active_plan_error=False + ) + self.assertEqual(result.status_code, 302) # Server plan status was reset self.server.refresh_from_db() @@ -538,6 +555,9 @@ class RemoteBillingAuthenticationTest(RemoteRealmBillingTestCase): def test_transfer_legacy_plan_scheduled_for_upgrade_from_server_to_realm( self, ) -> None: + self.login("desdemona") + desdemona = self.example_user("desdemona") + # Assert current server is not on any plan. self.assertIsNone(get_customer_by_remote_server(self.server)) @@ -580,6 +600,12 @@ class RemoteBillingAuthenticationTest(RemoteRealmBillingTestCase): self.add_mock_response() send_server_data_to_push_bouncer(consider_usage_statistics=False) + # Login to plan management. + result = self.execute_remote_billing_authentication_flow( + desdemona, server_on_active_plan_error=True + ) + self.assertEqual(result.status_code, 200) + # Server plan status stayed the same. self.server.refresh_from_db() self.assertEqual(self.server.plan_type, RemoteZulipServer.PLAN_TYPE_SELF_MANAGED_LEGACY) @@ -600,6 +626,12 @@ class RemoteBillingAuthenticationTest(RemoteRealmBillingTestCase): # Send server data to push bouncer. send_server_data_to_push_bouncer(consider_usage_statistics=False) + # Login to plan management. Performs customer migration from server to realms. + result = self.execute_remote_billing_authentication_flow( + desdemona, server_on_active_plan_error=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) @@ -635,6 +667,9 @@ class RemoteBillingAuthenticationTest(RemoteRealmBillingTestCase): def test_transfer_business_plan_from_server_to_realm( self, ) -> None: + self.login("desdemona") + desdemona = self.example_user("desdemona") + # Assert current server is not on any plan. self.assertIsNone(get_customer_by_remote_server(self.server)) self.assertEqual(self.server.plan_type, RemoteZulipServer.PLAN_TYPE_SELF_MANAGED) @@ -666,6 +701,12 @@ class RemoteBillingAuthenticationTest(RemoteRealmBillingTestCase): self.add_mock_response() send_server_data_to_push_bouncer(consider_usage_statistics=False) + # Login to plan management. + result = self.execute_remote_billing_authentication_flow( + desdemona, server_on_active_plan_error=True + ) + self.assertEqual(result.status_code, 200) + # Server plan status stayed the same. self.server.refresh_from_db() self.assertEqual(self.server.plan_type, RemoteZulipServer.PLAN_TYPE_BUSINESS) @@ -686,6 +727,12 @@ class RemoteBillingAuthenticationTest(RemoteRealmBillingTestCase): # Send server data to push bouncer. send_server_data_to_push_bouncer(consider_usage_statistics=False) + # Login to plan management. Performs customer migration from server to realms. + result = self.execute_remote_billing_authentication_flow( + desdemona, server_on_active_plan_error=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) @@ -709,6 +756,9 @@ class RemoteBillingAuthenticationTest(RemoteRealmBillingTestCase): @responses.activate def test_transfer_plan_from_server_to_realm_edge_cases(self) -> None: + self.login("desdemona") + desdemona = self.example_user("desdemona") + # CASE: Server has no customer self.assertIsNone(get_customer_by_remote_server(self.server)) @@ -716,6 +766,10 @@ class RemoteBillingAuthenticationTest(RemoteRealmBillingTestCase): self.add_mock_response() send_server_data_to_push_bouncer(consider_usage_statistics=False) + # Login to plan management. + result = self.execute_remote_billing_authentication_flow(desdemona) + self.assertEqual(result.status_code, 302) + # Still no customer. self.assertIsNone(get_customer_by_remote_server(self.server)) @@ -726,6 +780,12 @@ class RemoteBillingAuthenticationTest(RemoteRealmBillingTestCase): # Send server data to push bouncer. send_server_data_to_push_bouncer(consider_usage_statistics=False) + # Login to plan management. + result = self.execute_remote_billing_authentication_flow( + desdemona, first_time_login=False, expect_tos=False + ) + self.assertEqual(result.status_code, 302) + # Server still has no plan. self.assertIsNone(get_current_plan_by_customer(server_customer)) @@ -740,6 +800,12 @@ class RemoteBillingAuthenticationTest(RemoteRealmBillingTestCase): # Send server data to push bouncer. send_server_data_to_push_bouncer(consider_usage_statistics=False) + # Login to plan management. + result = self.execute_remote_billing_authentication_flow( + desdemona, server_on_active_plan_error=True + ) + self.assertEqual(result.status_code, 200) + # Server stays on the same plan. server_plan = get_current_plan_by_customer(server_customer) assert server_plan is not None @@ -753,8 +819,11 @@ class RemoteBillingAuthenticationTest(RemoteRealmBillingTestCase): self.server.plan_type = RemoteZulipServer.PLAN_TYPE_BUSINESS self.server.save(update_fields=["plan_type"]) - # Send server data to push bouncer. - send_server_data_to_push_bouncer(consider_usage_statistics=False) + # Login to plan management. + result = self.execute_remote_billing_authentication_flow( + desdemona, server_on_active_plan_error=True + ) + self.assertEqual(result.status_code, 200) # Server stays on the same plan. server_customer.refresh_from_db() diff --git a/corporate/tests/test_stripe.py b/corporate/tests/test_stripe.py index 3b9f2d9126..7ad184d7de 100644 --- a/corporate/tests/test_stripe.py +++ b/corporate/tests/test_stripe.py @@ -6070,11 +6070,20 @@ class TestRemoteRealmBillingFlow(StripeTestCase, RemoteRealmBillingTestCase): self.assertEqual(server_customer_plan.status, CustomerPlan.ACTIVE) self.assertEqual(remote_server.plan_type, RemoteZulipServer.PLAN_TYPE_SELF_MANAGED_LEGACY) - # Upload data. Performs customer migration from server to realms. + # Upload data. with time_machine.travel(self.now, tick=False): self.add_mock_response() send_server_data_to_push_bouncer(consider_usage_statistics=False) + self.login("hamlet") + hamlet = self.example_user("hamlet") + billing_base_url = self.billing_session.billing_base_url + + # Login. Performs customer migration from server to realms. + result = self.execute_remote_billing_authentication_flow(hamlet) + self.assertEqual(result.status_code, 302) + self.assertEqual(result["Location"], f"{billing_base_url}/plans/") + remote_server.refresh_from_db() server_customer_plan.refresh_from_db() self.assertEqual(server_customer_plan.status, CustomerPlan.ENDED) @@ -6092,15 +6101,6 @@ class TestRemoteRealmBillingFlow(StripeTestCase, RemoteRealmBillingTestCase): self.assertEqual(customer_plan.tier, CustomerPlan.TIER_SELF_HOSTED_LEGACY) self.assertEqual(customer_plan.status, CustomerPlan.ACTIVE) - self.login("hamlet") - hamlet = self.example_user("hamlet") - billing_base_url = self.billing_session.billing_base_url - - # Login - result = self.execute_remote_billing_authentication_flow(hamlet) - self.assertEqual(result.status_code, 302) - self.assertEqual(result["Location"], f"{billing_base_url}/plans/") - # upgrade to business plan with time_machine.travel(self.now, tick=False): result = self.client_get(f"{billing_base_url}/upgrade/", subdomain="selfhosting") diff --git a/corporate/views/remote_billing_page.py b/corporate/views/remote_billing_page.py index 84874c36e8..4a51761580 100644 --- a/corporate/views/remote_billing_page.py +++ b/corporate/views/remote_billing_page.py @@ -47,7 +47,11 @@ from zerver.lib.exceptions import ( RemoteBillingAuthenticationError, RemoteRealmServerMismatchError, ) -from zerver.lib.remote_server import RealmDataForAnalytics, UserDataForRemoteBilling +from zerver.lib.remote_server import ( + RealmDataForAnalytics, + UserDataForRemoteBilling, + get_realms_info_for_push_bouncer, +) from zerver.lib.response import json_success from zerver.lib.send_email import FromAddress, send_email from zerver.lib.timestamp import datetime_to_timestamp @@ -61,6 +65,7 @@ from zilencer.models import ( RemoteZulipServer, get_remote_server_by_uuid, ) +from zilencer.views import handle_customer_migration_from_server_to_realms billing_logger = logging.getLogger("corporate.stripe") @@ -194,6 +199,23 @@ def remote_realm_billing_finalize_login( # pretty recently. (And we generally don't delete these at all.) raise AssertionError + try: + handle_customer_migration_from_server_to_realms( + server=remote_server, realms=get_realms_info_for_push_bouncer() + ) + except Exception: # nocoverage + billing_logger.exception( + "%s: Failed to migrate customer from server (id: %s) to realms", + request.path, + remote_server.id, + stack_info=True, + ) + raise JsonableError( + _( + "Failed to migrate customer from server to realms. Please contact support for assistance." + ) + ) + # Redirect to error page if server is on an active plan server_customer = get_customer_by_remote_server(remote_server) if server_customer is not None: diff --git a/zilencer/views.py b/zilencer/views.py index 17728f5ded..46f5b9acb9 100644 --- a/zilencer/views.py +++ b/zilencer/views.py @@ -993,21 +993,6 @@ def remote_server_post_analytics( if remote_server_version_updated: fix_remote_realm_foreign_keys(server, realms) - try: - handle_customer_migration_from_server_to_realms(server, realms) - except Exception: # nocoverage - logger.exception( - "%s: Failed to migrate customer from server (id: %s) to realms", - request.path, - server.id, - stack_info=True, - ) - raise JsonableError( - _( - "Failed to migrate customer from server to realms. Please contact support for assistance." - ) - ) - realm_id_to_remote_realm = build_realm_id_to_remote_realm_dict(server, realms) remote_realm_counts = [