From 044cb820f8c2d1fb47fecf83b4e52280ac1d8cf9 Mon Sep 17 00:00:00 2001
From: Aman Agrawal <amanagr@zulip.com>
Date: Tue, 5 Dec 2023 06:39:21 +0000
Subject: [PATCH] stripe: Fix legacy server upgrade to business plan.

I had pushed a similar change in #28017 but seems to have been lost.
---
 corporate/lib/stripe.py               | 10 +++++++++-
 corporate/lib/stripe_event_handler.py | 13 +++++++------
 2 files changed, 16 insertions(+), 7 deletions(-)

diff --git a/corporate/lib/stripe.py b/corporate/lib/stripe.py
index c57bf54808..f90dcbb084 100644
--- a/corporate/lib/stripe.py
+++ b/corporate/lib/stripe.py
@@ -408,6 +408,14 @@ class UpgradeWithExistingPlanError(BillingError):
         )
 
 
+class InvalidPlanUpgradeError(BillingError):  # nocoverage
+    def __init__(self, message: str) -> None:
+        super().__init__(
+            "invalid plan upgrade",
+            message,
+        )
+
+
 class InvalidBillingScheduleError(Exception):
     def __init__(self, billing_schedule: int) -> None:
         self.message = f"Unknown billing_schedule: {billing_schedule}"
@@ -1067,7 +1075,7 @@ class BillingSession(ABC):
         assert plan is not None
         type_of_plan_change = self.get_type_of_plan_tier_change(plan.tier, new_plan_tier)
         if type_of_plan_change != PlanTierChangeType.UPGRADE:
-            raise BillingError(
+            raise InvalidPlanUpgradeError(
                 f"Cannot upgrade from {plan.name} to {CustomerPlan.name_from_tier(new_plan_tier)}"
             )
 
diff --git a/corporate/lib/stripe_event_handler.py b/corporate/lib/stripe_event_handler.py
index ceb571a8fa..da90cbe849 100644
--- a/corporate/lib/stripe_event_handler.py
+++ b/corporate/lib/stripe_event_handler.py
@@ -6,11 +6,11 @@ from django.conf import settings
 
 from corporate.lib.stripe import (
     BillingError,
+    InvalidPlanUpgradeError,
     RealmBillingSession,
     RemoteRealmBillingSession,
     RemoteServerBillingSession,
     UpgradeWithExistingPlanError,
-    ensure_customer_does_not_have_active_plan,
 )
 from corporate.models import Customer, CustomerPlan, Event, PaymentIntent, Session
 from zerver.models import get_active_user_profile_by_id_in_realm
@@ -123,10 +123,13 @@ def handle_payment_intent_succeeded_event(
         description=description,
         discountable=False,
     )
+    billing_session = get_billing_session_for_stripe_webhook(
+        payment_intent.customer, metadata.get("user_id")
+    )
     plan_tier = int(metadata["plan_tier"])
     try:
-        ensure_customer_does_not_have_active_plan(payment_intent.customer)
-    except UpgradeWithExistingPlanError as e:
+        billing_session.ensure_current_plan_is_upgradable(payment_intent.customer, plan_tier)
+    except (UpgradeWithExistingPlanError, InvalidPlanUpgradeError) as e:
         stripe_invoice = stripe.Invoice.create(
             auto_advance=True,
             collection_method="charge_automatically",
@@ -138,9 +141,6 @@ def handle_payment_intent_succeeded_event(
         stripe.Invoice.finalize_invoice(stripe_invoice)
         raise e
 
-    billing_session = get_billing_session_for_stripe_webhook(
-        payment_intent.customer, metadata.get("user_id")
-    )
     billing_session.process_initial_upgrade(
         plan_tier,
         int(metadata["licenses"]),
@@ -148,4 +148,5 @@ def handle_payment_intent_succeeded_event(
         int(metadata["billing_schedule"]),
         True,
         False,
+        billing_session.get_remote_server_legacy_plan(payment_intent.customer),
     )