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.
This commit is contained in:
Prakhar Pratyush 2024-01-05 18:54:15 +05:30 committed by Tim Abbott
parent 4ab9cd7cf2
commit 4715a058b0
4 changed files with 106 additions and 30 deletions

View File

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

View File

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

View File

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

View File

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