corporate: Move update_sponsorship_status to BillingSession class.

Moves `update_sponsorship_status` to BillingSession abstract class
as `update_customer_sponsorship_status`.

Updates the support views to have a helper for updating this on a
realm: `update_realm_sponsorship_status`.

analytics/tests/test_support_views.py

@@ -6,7 +6,8 @@ import orjson
 from django.utils.timezone import now as timezone_now
 from typing_extensions import override
 
-from corporate.lib.stripe import add_months, update_sponsorship_status
+from corporate.lib.stripe import add_months
+from corporate.lib.support import update_realm_sponsorship_status
 from corporate.models import Customer, CustomerPlan, LicenseLedger, get_customer_by_realm
 from zerver.actions.invites import do_create_multiuse_invite_link
 from zerver.actions.realm_settings import do_change_realm_org_type, do_send_realm_reactivation_email
@@ -558,8 +559,9 @@ class TestSupportEndpoint(ZulipTestCase):
         self.assertFalse(customer.sponsorship_pending)
 
     def test_approve_sponsorship(self) -> None:
+        support_admin = self.example_user("iago")
         lear_realm = get_realm("lear")
-        update_sponsorship_status(lear_realm, True, acting_user=None)
+        update_realm_sponsorship_status(lear_realm, True, acting_user=support_admin)
         king_user = self.lear_user("king")
         king_user.role = UserProfile.ROLE_REALM_OWNER
         king_user.save()

analytics/views/support.py

@@ -59,13 +59,13 @@ if settings.BILLING_ENABLED:
         make_end_of_cycle_updates_if_needed,
         switch_realm_from_standard_to_plus_plan,
         update_billing_method_of_current_plan,
-        update_sponsorship_status,
         void_all_open_invoices,
     )
     from corporate.lib.support import (
         approve_realm_sponsorship,
         attach_discount_to_realm,
         get_discount_for_realm,
+        update_realm_sponsorship_status,
     )
     from corporate.models import (
         Customer,
@@ -254,10 +254,10 @@ def support(
                 ] = f"Billing method of {realm.string_id} updated to charge automatically."
         elif sponsorship_pending is not None:
             if sponsorship_pending:
-                update_sponsorship_status(realm, True, acting_user=acting_user)
+                update_realm_sponsorship_status(realm, True, acting_user=acting_user)
                 context["success_message"] = f"{realm.string_id} marked as pending sponsorship."
             else:
-                update_sponsorship_status(realm, False, acting_user=acting_user)
+                update_realm_sponsorship_status(realm, False, acting_user=acting_user)
                 context["success_message"] = f"{realm.string_id} is no longer pending sponsorship."
         elif approve_sponsorship:
             approve_realm_sponsorship(realm, acting_user=acting_user)

corporate/lib/stripe.py

@@ -331,6 +331,7 @@ class AuditLogEventType(Enum):
     CUSTOMER_PLAN_CREATED = 3
     DISCOUNT_CHANGED = 4
     SPONSORSHIP_APPROVED = 5
+    SPONSORSHIP_PENDING_STATUS_CHANGED = 6
 
 
 class BillingSessionAuditLogEventError(Exception):
@@ -450,6 +451,18 @@ class BillingSession(ABC):
             extra_data={"old_discount": old_discount, "new_discount": discount},
         )
 
+    def update_customer_sponsorship_status(self, sponsorship_pending: bool) -> None:
+        customer = self.get_customer()
+        if customer is None:
+            customer = self.update_or_create_customer()
+        customer.sponsorship_pending = sponsorship_pending
+        customer.save(update_fields=["sponsorship_pending"])
+        self.write_to_audit_log(
+            event_type=AuditLogEventType.SPONSORSHIP_PENDING_STATUS_CHANGED,
+            event_time=timezone_now(),
+            extra_data={"sponsorship_pending": sponsorship_pending},
+        )
+
 
 class RealmBillingSession(BillingSession):
     def __init__(self, user: UserProfile, realm: Optional[Realm] = None) -> None:
@@ -478,6 +491,8 @@ class RealmBillingSession(BillingSession):
             return RealmAuditLog.REALM_DISCOUNT_CHANGED
         elif event_type is AuditLogEventType.SPONSORSHIP_APPROVED:
             return RealmAuditLog.REALM_SPONSORSHIP_APPROVED
+        elif event_type is AuditLogEventType.SPONSORSHIP_PENDING_STATUS_CHANGED:
+            return RealmAuditLog.REALM_SPONSORSHIP_PENDING_STATUS_CHANGED
         else:
             raise BillingSessionAuditLogEventError(event_type)
 
@@ -1124,21 +1139,6 @@ def is_realm_on_free_trial(realm: Realm) -> bool:
     return plan is not None and plan.is_free_trial()
 
 
-def update_sponsorship_status(
-    realm: Realm, sponsorship_pending: bool, *, acting_user: Optional[UserProfile]
-) -> None:
-    customer, _ = Customer.objects.get_or_create(realm=realm)
-    customer.sponsorship_pending = sponsorship_pending
-    customer.save(update_fields=["sponsorship_pending"])
-    RealmAuditLog.objects.create(
-        realm=realm,
-        acting_user=acting_user,
-        event_type=RealmAuditLog.REALM_SPONSORSHIP_PENDING_STATUS_CHANGED,
-        event_time=timezone_now(),
-        extra_data={"sponsorship_pending": sponsorship_pending},
-    )
-
-
 def is_sponsored_realm(realm: Realm) -> bool:
     return realm.plan_type == Realm.PLAN_TYPE_STANDARD_FREE
 

corporate/lib/support.py

@@ -34,3 +34,10 @@ def attach_discount_to_realm(realm: Realm, discount: Decimal, *, acting_user: Us
 def approve_realm_sponsorship(realm: Realm, *, acting_user: UserProfile) -> None:
     billing_session = RealmBillingSession(acting_user, realm)
     billing_session.approve_sponsorship()
+
+
+def update_realm_sponsorship_status(
+    realm: Realm, sponsorship_pending: bool, *, acting_user: UserProfile
+) -> None:
+    billing_session = RealmBillingSession(acting_user, realm)
+    billing_session.update_customer_sponsorship_status(sponsorship_pending)

corporate/tests/test_stripe.py

@@ -76,13 +76,13 @@ from corporate.lib.stripe import (
     update_license_ledger_for_automanaged_plan,
     update_license_ledger_for_manual_plan,
     update_license_ledger_if_needed,
-    update_sponsorship_status,
     void_all_open_invoices,
 )
 from corporate.lib.support import (
     approve_realm_sponsorship,
     attach_discount_to_realm,
     get_discount_for_realm,
+    update_realm_sponsorship_status,
 )
 from corporate.models import (
     Customer,
@@ -2475,21 +2475,6 @@ class StripeTest(StripeTestCase):
         # card on file, and should show it
         # TODO
 
-    def test_update_sponsorship_status(self) -> None:
-        lear = get_realm("lear")
-        iago = self.example_user("iago")
-        update_sponsorship_status(lear, True, acting_user=iago)
-        customer = get_customer_by_realm(realm=lear)
-        assert customer is not None
-        self.assertTrue(customer.sponsorship_pending)
-        realm_audit_log = RealmAuditLog.objects.filter(
-            event_type=RealmAuditLog.REALM_SPONSORSHIP_PENDING_STATUS_CHANGED
-        ).last()
-        assert realm_audit_log is not None
-        expected_extra_data = {"sponsorship_pending": True}
-        self.assertEqual(realm_audit_log.extra_data, expected_extra_data)
-        self.assertEqual(realm_audit_log.acting_user, iago)
-
     @mock_stripe()
     def test_replace_payment_method(self, *mocks: Mock) -> None:
         user = self.example_user("hamlet")
@@ -5062,3 +5047,18 @@ class TestSupportBillingHelpers(StripeTestCase):
         self.assertEqual(message.content, expected_message)
         self.assertEqual(message.recipient.type, Recipient.PERSONAL)
         self.assertEqual(message.recipient_id, recipient_id)
+
+    def test_update_realm_sponsorship_status(self) -> None:
+        lear = get_realm("lear")
+        iago = self.example_user("iago")
+        update_realm_sponsorship_status(lear, True, acting_user=iago)
+        customer = get_customer_by_realm(realm=lear)
+        assert customer is not None
+        self.assertTrue(customer.sponsorship_pending)
+        realm_audit_log = RealmAuditLog.objects.filter(
+            event_type=RealmAuditLog.REALM_SPONSORSHIP_PENDING_STATUS_CHANGED
+        ).last()
+        assert realm_audit_log is not None
+        expected_extra_data = {"sponsorship_pending": True}
+        self.assertEqual(realm_audit_log.extra_data, expected_extra_data)
+        self.assertEqual(realm_audit_log.acting_user, iago)

corporate/views/upgrade.py

@@ -24,7 +24,6 @@ from corporate.lib.stripe import (
     process_initial_upgrade,
     sign_string,
     unsign_string,
-    update_sponsorship_status,
     validate_licenses,
 )
 from corporate.lib.support import get_support_url
@@ -325,6 +324,7 @@ def sponsorship(
     description: str = REQ(),
 ) -> HttpResponse:
     realm = user.realm
+    billing_session = RealmBillingSession(user)
 
     requested_by = user.full_name
     user_role = user.get_role_name()
@@ -353,7 +353,7 @@ def sponsorship(
                 realm.org_type = org_type
                 realm.save(update_fields=["org_type"])
 
-            update_sponsorship_status(realm, True, acting_user=user)
+            billing_session.update_customer_sponsorship_status(True)
             do_make_user_billing_admin(user)
 
             org_type_display_name = get_org_type_display_name(org_type)