From 20d2e00cf9571a4783793ea3098b4b02f31a9960 Mon Sep 17 00:00:00 2001
From: Vishnu KS <yo@vishnuks.com>
Date: Wed, 30 Dec 2020 23:26:57 +0530
Subject: [PATCH] billing: Create update_license_ledger_for_manual_plan.

---
 corporate/lib/stripe.py        | 24 ++++++++++++++
 corporate/tests/test_stripe.py | 60 ++++++++++++++++++++++++++++++++++
 2 files changed, 84 insertions(+)

diff --git a/corporate/lib/stripe.py b/corporate/lib/stripe.py
index 8649b088de..4c3368301a 100644
--- a/corporate/lib/stripe.py
+++ b/corporate/lib/stripe.py
@@ -629,6 +629,30 @@ def process_initial_upgrade(
     do_change_plan_type(realm, Realm.STANDARD, acting_user=user)
 
 
+def update_license_ledger_for_manual_plan(
+    plan: CustomerPlan,
+    event_time: datetime,
+    licenses: Optional[int] = None,
+    licenses_at_next_renewal: Optional[int] = None,
+) -> None:
+    if licenses is not None:
+        assert get_latest_seat_count(plan.customer.realm) <= licenses
+        assert licenses > plan.licenses()
+        LicenseLedger.objects.create(
+            plan=plan, event_time=event_time, licenses=licenses, licenses_at_next_renewal=licenses
+        )
+    elif licenses_at_next_renewal is not None:
+        assert get_latest_seat_count(plan.customer.realm) <= licenses_at_next_renewal
+        LicenseLedger.objects.create(
+            plan=plan,
+            event_time=event_time,
+            licenses=plan.licenses(),
+            licenses_at_next_renewal=licenses_at_next_renewal,
+        )
+    else:
+        raise AssertionError("Pass licenses or licenses_at_next_renewal")
+
+
 def update_license_ledger_for_automanaged_plan(
     realm: Realm, plan: CustomerPlan, event_time: datetime
 ) -> None:
diff --git a/corporate/tests/test_stripe.py b/corporate/tests/test_stripe.py
index 9f89f4fb06..5798f481b6 100644
--- a/corporate/tests/test_stripe.py
+++ b/corporate/tests/test_stripe.py
@@ -42,6 +42,7 @@ from corporate.lib.stripe import (
     unsign_string,
     update_billing_method_of_current_plan,
     update_license_ledger_for_automanaged_plan,
+    update_license_ledger_for_manual_plan,
     update_license_ledger_if_needed,
     update_or_create_stripe_customer,
     update_sponsorship_status,
@@ -2882,6 +2883,65 @@ class LicenseLedgerTest(StripeTestCase):
             ],
         )
 
+    def test_update_license_ledger_for_manual_plan(self) -> None:
+        realm = get_realm("zulip")
+
+        with patch("corporate.lib.stripe.timezone_now", return_value=self.now):
+            self.local_upgrade(self.seat_count + 1, False, CustomerPlan.ANNUAL, "token")
+
+        plan = get_current_plan_by_realm(realm)
+        assert plan is not None
+
+        with patch("corporate.lib.stripe.get_latest_seat_count", return_value=self.seat_count):
+            update_license_ledger_for_manual_plan(plan, self.now, licenses=self.seat_count + 3)
+            self.assertEqual(plan.licenses(), self.seat_count + 3)
+            self.assertEqual(plan.licenses_at_next_renewal(), self.seat_count + 3)
+
+        with patch("corporate.lib.stripe.get_latest_seat_count", return_value=self.seat_count):
+            with self.assertRaises(AssertionError):
+                update_license_ledger_for_manual_plan(plan, self.now, licenses=self.seat_count)
+
+        with patch("corporate.lib.stripe.get_latest_seat_count", return_value=self.seat_count):
+            update_license_ledger_for_manual_plan(
+                plan, self.now, licenses_at_next_renewal=self.seat_count
+            )
+            self.assertEqual(plan.licenses(), self.seat_count + 3)
+            self.assertEqual(plan.licenses_at_next_renewal(), self.seat_count)
+
+        with patch("corporate.lib.stripe.get_latest_seat_count", return_value=self.seat_count):
+            with self.assertRaises(AssertionError):
+                update_license_ledger_for_manual_plan(
+                    plan, self.now, licenses_at_next_renewal=self.seat_count - 1
+                )
+
+        with patch("corporate.lib.stripe.get_latest_seat_count", return_value=self.seat_count):
+            update_license_ledger_for_manual_plan(plan, self.now, licenses=self.seat_count + 10)
+            self.assertEqual(plan.licenses(), self.seat_count + 10)
+            self.assertEqual(plan.licenses_at_next_renewal(), self.seat_count + 10)
+
+        make_end_of_cycle_updates_if_needed(plan, self.next_year)
+        self.assertEqual(plan.licenses(), self.seat_count + 10)
+
+        ledger_entries = list(
+            LicenseLedger.objects.values_list(
+                "is_renewal", "event_time", "licenses", "licenses_at_next_renewal"
+            ).order_by("id")
+        )
+
+        self.assertEqual(
+            ledger_entries,
+            [
+                (True, self.now, self.seat_count + 1, self.seat_count + 1),
+                (False, self.now, self.seat_count + 3, self.seat_count + 3),
+                (False, self.now, self.seat_count + 3, self.seat_count),
+                (False, self.now, self.seat_count + 10, self.seat_count + 10),
+                (True, self.next_year, self.seat_count + 10, self.seat_count + 10),
+            ],
+        )
+
+        with self.assertRaises(AssertionError):
+            update_license_ledger_for_manual_plan(plan, self.now)
+
     def test_user_changes(self) -> None:
         self.local_upgrade(self.seat_count, True, CustomerPlan.ANNUAL, "token")
         user = do_create_user("email", "password", get_realm("zulip"), "name", acting_user=None)