From 5d07666362efbf5edebf04ea5ca6c8208b922385 Mon Sep 17 00:00:00 2001 From: Lauryn Menard Date: Thu, 2 Nov 2023 18:17:08 +0100 Subject: [PATCH] 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 +++-- analytics/views/support.py | 6 ++--- corporate/lib/stripe.py | 30 ++++++++++++------------- corporate/lib/support.py | 7 ++++++ corporate/tests/test_stripe.py | 32 +++++++++++++-------------- corporate/views/upgrade.py | 4 ++-- 6 files changed, 47 insertions(+), 38 deletions(-) diff --git a/analytics/tests/test_support_views.py b/analytics/tests/test_support_views.py index 6ffa337965..abbea473e6 100644 --- a/analytics/tests/test_support_views.py +++ b/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() diff --git a/analytics/views/support.py b/analytics/views/support.py index 294423096c..2118aad66d 100644 --- a/analytics/views/support.py +++ b/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) diff --git a/corporate/lib/stripe.py b/corporate/lib/stripe.py index 24f6745dc5..37f41e7513 100644 --- a/corporate/lib/stripe.py +++ b/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 diff --git a/corporate/lib/support.py b/corporate/lib/support.py index 34e8923485..de36dae11a 100644 --- a/corporate/lib/support.py +++ b/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) diff --git a/corporate/tests/test_stripe.py b/corporate/tests/test_stripe.py index 12d3f433e4..383222f8a9 100644 --- a/corporate/tests/test_stripe.py +++ b/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) diff --git a/corporate/views/upgrade.py b/corporate/views/upgrade.py index 4165004798..0f87d9b46b 100644 --- a/corporate/views/upgrade.py +++ b/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)