From e65f3cf6576005b036afe7ff0ad5a2484f52d2d8 Mon Sep 17 00:00:00 2001 From: Lauryn Menard Date: Tue, 31 Dec 2024 15:13:41 +0100 Subject: [PATCH] corporate: Create license ledger for automanaged plan migrations. If we move a paid plan from a remote server to a remote realm, and the plan has automated license management, then we create an updated license ledger entry when we move the plan for the remote realm billing data so that we have an accurate user count for licenses when the plan is next invoiced. --- corporate/tests/test_remote_billing.py | 20 ++++++++++++++++ zilencer/views.py | 32 +++++++++++++++++++++++--- 2 files changed, 49 insertions(+), 3 deletions(-) diff --git a/corporate/tests/test_remote_billing.py b/corporate/tests/test_remote_billing.py index 3a2b8e162b..8bb1bea39b 100644 --- a/corporate/tests/test_remote_billing.py +++ b/corporate/tests/test_remote_billing.py @@ -17,6 +17,7 @@ from corporate.lib.remote_billing_util import ( from corporate.lib.stripe import RemoteRealmBillingSession, RemoteServerBillingSession, add_months from corporate.models import ( CustomerPlan, + LicenseLedger, get_current_plan_by_customer, get_customer_by_remote_realm, get_customer_by_remote_server, @@ -1019,6 +1020,15 @@ class RemoteBillingAuthenticationTest(RemoteRealmBillingTestCase): billing_schedule=CustomerPlan.BILLING_SCHEDULE_ANNUAL, tier=CustomerPlan.TIER_SELF_HOSTED_BUSINESS, status=CustomerPlan.ACTIVE, + automanage_licenses=True, + ) + initial_license_count = 100 + LicenseLedger.objects.create( + plan=server_plan, + is_renewal=True, + event_time=timezone_now(), + licenses=initial_license_count, + licenses_at_next_renewal=initial_license_count, ) self.server.plan_type = RemoteZulipServer.PLAN_TYPE_BUSINESS self.server.save(update_fields=["plan_type"]) @@ -1096,6 +1106,16 @@ class RemoteBillingAuthenticationTest(RemoteRealmBillingTestCase): self.assertEqual(plan.tier, CustomerPlan.TIER_SELF_HOSTED_BUSINESS) self.assertEqual(plan.status, CustomerPlan.ACTIVE) + # Check that an updated license ledger entry was created. + billing_session = RemoteRealmBillingSession(remote_realm=remote_realm_with_plan) + license_ledger = billing_session.get_last_ledger_for_automanaged_plan_if_exists() + billable_licenses = billing_session.get_billable_licenses_for_customer(customer, plan.tier) + assert license_ledger is not None + self.assertNotEqual(initial_license_count, billable_licenses) + self.assertEqual(license_ledger.licenses, initial_license_count) + self.assertEqual(license_ledger.licenses_at_next_renewal, billable_licenses) + self.assertFalse(license_ledger.is_renewal) + @responses.activate def test_transfer_plan_from_server_to_realm_edge_cases(self) -> None: self.login("desdemona") diff --git a/zilencer/views.py b/zilencer/views.py index 642d8bf905..13e99698f5 100644 --- a/zilencer/views.py +++ b/zilencer/views.py @@ -74,6 +74,7 @@ from zerver.models.realm_audit_logs import AuditLogEventType from zerver.models.realms import DisposableEmailError from zerver.views.push_notifications import validate_token from zilencer.auth import InvalidZulipServerKeyError +from zilencer.lib.remote_counts import MissingDataError from zilencer.models import ( RemoteInstallationCount, RemotePushDeviceToken, @@ -1039,7 +1040,7 @@ def get_human_user_realm_uuids( def handle_customer_migration_from_server_to_realm( server: RemoteZulipServer, ) -> None: - from corporate.lib.stripe import RemoteServerBillingSession + from corporate.lib.stripe import RemoteRealmBillingSession, RemoteServerBillingSession server_billing_session = RemoteServerBillingSession(server) server_customer = server_billing_session.get_customer() @@ -1050,7 +1051,7 @@ def handle_customer_migration_from_server_to_realm( # If we have a pending sponsorship request, defer moving any # data until the sponsorship request has been processed. This # avoids a race where a sponsorship request made at the server - # level gets approved after the legacy plan has already been + # level gets approved after the active plan has already been # moved to the sole human RemoteRealm, which would violate # invariants. return @@ -1058,7 +1059,7 @@ def handle_customer_migration_from_server_to_realm( server_plan = get_current_plan_by_customer(server_customer) if server_plan is None: # If the server has no current plan, either because it never - # had one or because a previous legacy plan was migrated to + # had one or because a previous active plan was migrated to # the RemoteRealm object, there's nothing to potentially # migrate. return @@ -1132,6 +1133,31 @@ def handle_customer_migration_from_server_to_realm( ).format(support_email=FromAddress.SUPPORT) ) + # We successfully moved the plan from the remote server to the remote realm. + # Update the license ledger for paid plans with automated license management. + remote_realm_customer = get_customer_by_remote_realm(remote_realm) + assert remote_realm_customer is not None + moved_customer_plan = get_current_plan_by_customer(remote_realm_customer) + assert moved_customer_plan is not None + if moved_customer_plan.is_a_paid_plan() and moved_customer_plan.automanage_licenses: + remote_realm_billing_session = RemoteRealmBillingSession(remote_realm=remote_realm) + try: + remote_realm_billing_session.update_license_ledger_for_automanaged_plan( + moved_customer_plan, event_time + ) + except MissingDataError: # nocoverage + logger.warning( + "Failed to migrate customer from server (id: %s) to realm (id: %s): RemoteZulipServer has stale " + "audit log data and cannot update license ledger for plan with automated license management.", + 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"])