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

@@ -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.
+        # Login to plan management. Performs customer migration from server to realms.
-        send_server_data_to_push_bouncer(consider_usage_statistics=False)
+        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.
+        # Login to plan management.
-        send_server_data_to_push_bouncer(consider_usage_statistics=False)
+        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()

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

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:

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