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 = [