test_classes: Create helper function for subscribing realm to manual plan.

corporate/tests/test_stripe.py

@@ -3258,3 +3258,47 @@ class InvoiceTest(StripeTestCase):
         invoice_plans_as_needed(self.next_month)
         plan = CustomerPlan.objects.first()
         self.assertEqual(plan.next_invoice_date, self.next_month + timedelta(days=29))
+
+
+class TestTestClasses(ZulipTestCase):
+    def test_subscribe_realm_to_manual_license_management_plan(self) -> None:
+        realm = get_realm("zulip")
+        plan, ledger = self.subscribe_realm_to_manual_license_management_plan(
+            realm, 50, 60, CustomerPlan.ANNUAL
+        )
+
+        plan.refresh_from_db()
+        self.assertEqual(plan.automanage_licenses, False)
+        self.assertEqual(plan.billing_schedule, CustomerPlan.ANNUAL)
+        self.assertEqual(plan.tier, CustomerPlan.STANDARD)
+        self.assertEqual(plan.licenses(), 50)
+        self.assertEqual(plan.licenses_at_next_renewal(), 60)
+
+        ledger.refresh_from_db()
+        self.assertEqual(ledger.plan, plan)
+        self.assertEqual(ledger.licenses, 50)
+        self.assertEqual(ledger.licenses_at_next_renewal, 60)
+
+        realm.refresh_from_db()
+        self.assertEqual(realm.plan_type, Realm.STANDARD)
+
+    def test_subscribe_realm_to_monthly_plan_on_manual_license_management(self) -> None:
+        realm = get_realm("zulip")
+        plan, ledger = self.subscribe_realm_to_monthly_plan_on_manual_license_management(
+            realm, 20, 30
+        )
+
+        plan.refresh_from_db()
+        self.assertEqual(plan.automanage_licenses, False)
+        self.assertEqual(plan.billing_schedule, CustomerPlan.MONTHLY)
+        self.assertEqual(plan.tier, CustomerPlan.STANDARD)
+        self.assertEqual(plan.licenses(), 20)
+        self.assertEqual(plan.licenses_at_next_renewal(), 30)
+
+        ledger.refresh_from_db()
+        self.assertEqual(ledger.plan, plan)
+        self.assertEqual(ledger.licenses, 20)
+        self.assertEqual(ledger.licenses_at_next_renewal, 30)
+
+        realm.refresh_from_db()
+        self.assertEqual(realm.plan_type, Realm.STANDARD)

zerver/lib/test_classes.py

@@ -43,6 +43,7 @@ from django.utils.timezone import now as timezone_now
 from fakeldap import MockLDAP
 from two_factor.models import PhoneDevice
 
+from corporate.models import Customer, CustomerPlan, LicenseLedger
 from zerver.decorator import do_two_factor_login
 from zerver.lib.actions import (
     bulk_add_subscriptions,
@@ -1266,6 +1267,35 @@ Output:
         self.assertTrue(validation_func(new_member_user))
         self.assertFalse(validation_func(guest_user))
 
+    def subscribe_realm_to_manual_license_management_plan(
+        self, realm: Realm, licenses: int, licenses_at_next_renewal: int, billing_schedule: int
+    ) -> Tuple[CustomerPlan, LicenseLedger]:
+        customer, _ = Customer.objects.get_or_create(realm=realm)
+        plan = CustomerPlan.objects.create(
+            customer=customer,
+            automanage_licenses=False,
+            billing_cycle_anchor=timezone_now(),
+            billing_schedule=billing_schedule,
+            tier=CustomerPlan.STANDARD,
+        )
+        ledger = LicenseLedger.objects.create(
+            plan=plan,
+            is_renewal=True,
+            event_time=timezone_now(),
+            licenses=licenses,
+            licenses_at_next_renewal=licenses_at_next_renewal,
+        )
+        realm.plan_type = Realm.STANDARD
+        realm.save(update_fields=["plan_type"])
+        return plan, ledger
+
+    def subscribe_realm_to_monthly_plan_on_manual_license_management(
+        self, realm: Realm, licenses: int, licenses_at_next_renewal: int
+    ) -> Tuple[CustomerPlan, LicenseLedger]:
+        return self.subscribe_realm_to_manual_license_management_plan(
+            realm, licenses, licenses_at_next_renewal, CustomerPlan.MONTHLY
+        )
+
     @contextmanager
     def tornado_redirected_to_list(
         self, lst: List[Mapping[str, Any]], expected_num_events: int