stripe: Move process_initial_upgrade to BillingSession.

Moves the 'process_initial_upgrade' function to the
'BillingSession' abstract class.

This refactoring will help in minimizing duplicate code while
supporting both realm and remote_server customers.

corporate/lib/stripe.py

@@ -604,6 +604,104 @@ class BillingSession(ABC):
         )
         return stripe_session
 
+    # Only used for cloud signups
+    @catch_stripe_errors
+    def process_initial_upgrade(
+        self,
+        plan_tier: int,
+        licenses: int,
+        automanage_licenses: bool,
+        billing_schedule: int,
+        charge_automatically: bool,
+        free_trial: bool,
+    ) -> None:
+        customer = self.update_or_create_stripe_customer()
+        assert customer.stripe_customer_id is not None  # for mypy
+        ensure_customer_does_not_have_active_plan(customer)
+        (
+            billing_cycle_anchor,
+            next_invoice_date,
+            period_end,
+            price_per_license,
+        ) = compute_plan_parameters(
+            plan_tier,
+            automanage_licenses,
+            billing_schedule,
+            customer.default_discount,
+            free_trial,
+        )
+
+        # TODO: The correctness of this relies on user creation, deactivation, etc being
+        # in a transaction.atomic() with the relevant RealmAuditLog entries
+        with transaction.atomic():
+            if customer.exempt_from_license_number_check:
+                billed_licenses = licenses
+            else:
+                # billed_licenses can be greater than licenses if users are added between the start of
+                # this function (process_initial_upgrade) and now
+                current_licenses_count = self.current_count_for_billed_licenses()
+                billed_licenses = max(current_licenses_count, licenses)
+            plan_params = {
+                "automanage_licenses": automanage_licenses,
+                "charge_automatically": charge_automatically,
+                "price_per_license": price_per_license,
+                "discount": customer.default_discount,
+                "billing_cycle_anchor": billing_cycle_anchor,
+                "billing_schedule": billing_schedule,
+                "tier": plan_tier,
+            }
+            if free_trial:
+                plan_params["status"] = CustomerPlan.FREE_TRIAL
+            plan = CustomerPlan.objects.create(
+                customer=customer, next_invoice_date=next_invoice_date, **plan_params
+            )
+            ledger_entry = LicenseLedger.objects.create(
+                plan=plan,
+                is_renewal=True,
+                event_time=billing_cycle_anchor,
+                licenses=billed_licenses,
+                licenses_at_next_renewal=billed_licenses,
+            )
+            plan.invoiced_through = ledger_entry
+            plan.save(update_fields=["invoiced_through"])
+            self.write_to_audit_log(
+                event_type=AuditLogEventType.CUSTOMER_PLAN_CREATED,
+                event_time=billing_cycle_anchor,
+                extra_data=plan_params,
+            )
+
+        if not free_trial:
+            stripe.InvoiceItem.create(
+                currency="usd",
+                customer=customer.stripe_customer_id,
+                description=plan.name,
+                discountable=False,
+                period={
+                    "start": datetime_to_timestamp(billing_cycle_anchor),
+                    "end": datetime_to_timestamp(period_end),
+                },
+                quantity=billed_licenses,
+                unit_amount=price_per_license,
+            )
+
+            if charge_automatically:
+                collection_method = "charge_automatically"
+                days_until_due = None
+            else:
+                collection_method = "send_invoice"
+                days_until_due = DEFAULT_INVOICE_DAYS_UNTIL_DUE
+
+            stripe_invoice = stripe.Invoice.create(
+                auto_advance=True,
+                collection_method=collection_method,
+                customer=customer.stripe_customer_id,
+                days_until_due=days_until_due,
+                statement_descriptor=plan.name,
+            )
+            stripe.Invoice.finalize_invoice(stripe_invoice)
+
+        self.do_change_plan_type(tier=plan_tier)
+
 
 class RealmBillingSession(BillingSession):
     def __init__(self, user: UserProfile, realm: Optional[Realm] = None) -> None:
@@ -1055,106 +1153,6 @@ def do_deactivate_remote_server(remote_server: RemoteZulipServer) -> None:
     )
 
 
-# Only used for cloud signups
-@catch_stripe_errors
-def process_initial_upgrade(
-    user: UserProfile,
-    plan_tier: int,
-    licenses: int,
-    automanage_licenses: bool,
-    billing_schedule: int,
-    charge_automatically: bool,
-    free_trial: bool,
-) -> None:
-    billing_session = RealmBillingSession(user)
-    customer = billing_session.update_or_create_stripe_customer()
-    assert customer.stripe_customer_id is not None  # for mypy
-    ensure_customer_does_not_have_active_plan(customer)
-    (
-        billing_cycle_anchor,
-        next_invoice_date,
-        period_end,
-        price_per_license,
-    ) = compute_plan_parameters(
-        plan_tier,
-        automanage_licenses,
-        billing_schedule,
-        customer.default_discount,
-        free_trial,
-    )
-
-    # TODO: The correctness of this relies on user creation, deactivation, etc being
-    # in a transaction.atomic() with the relevant RealmAuditLog entries
-    with transaction.atomic():
-        if customer.exempt_from_license_number_check:
-            billed_licenses = licenses
-        else:
-            # billed_licenses can be greater than licenses if users are added between the start of
-            # this function (process_initial_upgrade) and now
-            current_licenses_count = billing_session.current_count_for_billed_licenses()
-            billed_licenses = max(current_licenses_count, licenses)
-        plan_params = {
-            "automanage_licenses": automanage_licenses,
-            "charge_automatically": charge_automatically,
-            "price_per_license": price_per_license,
-            "discount": customer.default_discount,
-            "billing_cycle_anchor": billing_cycle_anchor,
-            "billing_schedule": billing_schedule,
-            "tier": plan_tier,
-        }
-        if free_trial:
-            plan_params["status"] = CustomerPlan.FREE_TRIAL
-        plan = CustomerPlan.objects.create(
-            customer=customer, next_invoice_date=next_invoice_date, **plan_params
-        )
-        ledger_entry = LicenseLedger.objects.create(
-            plan=plan,
-            is_renewal=True,
-            event_time=billing_cycle_anchor,
-            licenses=billed_licenses,
-            licenses_at_next_renewal=billed_licenses,
-        )
-        plan.invoiced_through = ledger_entry
-        plan.save(update_fields=["invoiced_through"])
-        billing_session.write_to_audit_log(
-            event_type=AuditLogEventType.CUSTOMER_PLAN_CREATED,
-            event_time=billing_cycle_anchor,
-            extra_data=plan_params,
-        )
-
-    if not free_trial:
-        stripe.InvoiceItem.create(
-            currency="usd",
-            customer=customer.stripe_customer_id,
-            description=plan.name,
-            discountable=False,
-            period={
-                "start": datetime_to_timestamp(billing_cycle_anchor),
-                "end": datetime_to_timestamp(period_end),
-            },
-            quantity=billed_licenses,
-            unit_amount=price_per_license,
-        )
-
-        if charge_automatically:
-            collection_method = "charge_automatically"
-            days_until_due = None
-        else:
-            collection_method = "send_invoice"
-            days_until_due = DEFAULT_INVOICE_DAYS_UNTIL_DUE
-
-        stripe_invoice = stripe.Invoice.create(
-            auto_advance=True,
-            collection_method=collection_method,
-            customer=customer.stripe_customer_id,
-            days_until_due=days_until_due,
-            statement_descriptor=plan.name,
-        )
-        stripe.Invoice.finalize_invoice(stripe_invoice)
-
-    billing_session.do_change_plan_type(tier=plan_tier)
-
-
 def update_license_ledger_for_manual_plan(
     plan: CustomerPlan,
     event_time: datetime,

corporate/lib/stripe_event_handler.py

@@ -10,7 +10,6 @@ from corporate.lib.stripe import (
     RealmBillingSession,
     UpgradeWithExistingPlanError,
     ensure_customer_does_not_have_active_plan,
-    process_initial_upgrade,
 )
 from corporate.models import CustomerPlan, Event, PaymentIntent, Session
 from zerver.models import get_active_user_profile_by_id_in_realm
@@ -100,8 +99,7 @@ def handle_checkout_session_completed_event(
     ]:
         ensure_customer_does_not_have_active_plan(session.customer)
         billing_session.update_or_create_stripe_customer(payment_method)
-        process_initial_upgrade(
-            user,
+        billing_session.process_initial_upgrade(
             CustomerPlan.STANDARD,
             int(stripe_session.metadata["licenses"]),
             stripe_session.metadata["license_management"] == "automatic",
@@ -150,8 +148,8 @@ def handle_payment_intent_succeeded_event(
         stripe.Invoice.finalize_invoice(stripe_invoice)
         raise e
 
-    process_initial_upgrade(
-        user,
+    billing_session = RealmBillingSession(user)
+    billing_session.process_initial_upgrade(
         CustomerPlan.STANDARD,
         int(metadata["licenses"]),
         metadata["license_management"] == "automatic",

corporate/tests/test_stripe.py

@@ -64,7 +64,6 @@ from corporate.lib.stripe import (
     is_sponsored_realm,
     make_end_of_cycle_updates_if_needed,
     next_month,
-    process_initial_upgrade,
     sign_string,
     stripe_customer_has_credit_card_as_default_payment_method,
     stripe_get_customer,
@@ -625,8 +624,9 @@ class StripeTestCase(ZulipTestCase):
             free_trial: bool,
             *mock_args: Any,
         ) -> Any:
-            return process_initial_upgrade(
-                self.example_user("hamlet"),
+            hamlet = self.example_user("hamlet")
+            billing_session = RealmBillingSession(hamlet)
+            return billing_session.process_initial_upgrade(
                 CustomerPlan.STANDARD,
                 licenses,
                 automanage_licenses,
@@ -1980,7 +1980,7 @@ class StripeTest(StripeTestCase):
             else:
                 del_args = []
                 upgrade_params["licenses"] = licenses
-            with patch("corporate.views.upgrade.process_initial_upgrade"):
+            with patch("corporate.lib.stripe.BillingSession.process_initial_upgrade"):
                 stripe_session = stripe.checkout.Session()
                 stripe_session.id = "stripe_session_id"
                 stripe_session.url = "stripe_session_url"
@@ -2087,7 +2087,7 @@ class StripeTest(StripeTestCase):
         self.login_user(hamlet)
 
         with patch(
-            "corporate.lib.stripe_event_handler.process_initial_upgrade", side_effect=Exception
+            "corporate.lib.stripe.BillingSession.process_initial_upgrade", side_effect=Exception
         ), self.assertLogs("corporate.stripe", "WARNING"):
             response = self.upgrade()
 

corporate/views/upgrade.py

@@ -18,7 +18,6 @@ from corporate.lib.stripe import (
     ensure_customer_does_not_have_active_plan,
     get_latest_seat_count,
     is_free_trial_offer_enabled,
-    process_initial_upgrade,
     sign_string,
     unsign_string,
     validate_licenses,
@@ -119,8 +118,8 @@ def upgrade(
         billing_schedule = {"annual": CustomerPlan.ANNUAL, "monthly": CustomerPlan.MONTHLY}[
             schedule
         ]
+        billing_session = RealmBillingSession(user)
         if charge_automatically:
-            billing_session = RealmBillingSession(user)
             stripe_checkout_session = (
                 billing_session.setup_upgrade_checkout_session_and_payment_intent(
                     CustomerPlan.STANDARD,
@@ -140,8 +139,7 @@ def upgrade(
                 },
             )
         else:
-            process_initial_upgrade(
-                user,
+            billing_session.process_initial_upgrade(
                 CustomerPlan.STANDARD,
                 licenses,
                 automanage_licenses,