diff --git a/corporate/lib/registration.py b/corporate/lib/registration.py index cd4ba6754a..248770e981 100644 --- a/corporate/lib/registration.py +++ b/corporate/lib/registration.py @@ -2,16 +2,23 @@ from django.conf import settings from django.utils.translation import gettext as _ from corporate.lib.stripe import LicenseLimitError, get_latest_seat_count, get_seat_count -from corporate.models import get_current_plan_by_realm +from corporate.models import CustomerPlan, get_current_plan_by_realm from zerver.actions.create_user import send_group_direct_message_to_admins from zerver.lib.exceptions import InvitationError from zerver.models import Realm, UserProfile from zerver.models.users import get_system_bot -def generate_licenses_low_warning_message_if_required(realm: Realm) -> str | None: +def get_plan_if_manual_license_management_enforced(realm: Realm) -> CustomerPlan | None: plan = get_current_plan_by_realm(realm) - if plan is None or plan.automanage_licenses: + if plan is None or plan.automanage_licenses or plan.customer.exempt_from_license_number_check: + return None + return plan + + +def generate_licenses_low_warning_message_if_required(realm: Realm) -> str | None: + plan = get_plan_if_manual_license_management_enforced(realm) + if plan is None: return None licenses_remaining = plan.licenses() - get_latest_seat_count(realm) @@ -67,13 +74,9 @@ def send_user_unable_to_signup_group_direct_message_to_admins( ) -def check_spare_licenses_available_for_adding_new_users( - realm: Realm, extra_non_guests_count: int = 0, extra_guests_count: int = 0 +def check_spare_licenses_available( + realm: Realm, plan: CustomerPlan, extra_non_guests_count: int = 0, extra_guests_count: int = 0 ) -> None: - plan = get_current_plan_by_realm(realm) - if plan is None or plan.automanage_licenses or plan.customer.exempt_from_license_number_check: - return - if plan.licenses() < get_seat_count( realm, extra_non_guests_count=extra_non_guests_count, extra_guests_count=extra_guests_count ): @@ -85,11 +88,15 @@ def check_spare_licenses_available_for_registering_new_user( user_email_to_add: str, role: int, ) -> None: + plan = get_plan_if_manual_license_management_enforced(realm) + if plan is None: + return + try: if role == UserProfile.ROLE_GUEST: - check_spare_licenses_available_for_adding_new_users(realm, extra_guests_count=1) + check_spare_licenses_available(realm, plan, extra_guests_count=1) else: - check_spare_licenses_available_for_adding_new_users(realm, extra_non_guests_count=1) + check_spare_licenses_available(realm, plan, extra_non_guests_count=1) except LicenseLimitError: send_user_unable_to_signup_group_direct_message_to_admins(realm, user_email_to_add) raise @@ -99,10 +106,12 @@ def check_spare_licenses_available_for_inviting_new_users( realm: Realm, extra_non_guests_count: int = 0, extra_guests_count: int = 0 ) -> None: num_invites = extra_non_guests_count + extra_guests_count + plan = get_plan_if_manual_license_management_enforced(realm) + if plan is None: + return + try: - check_spare_licenses_available_for_adding_new_users( - realm, extra_non_guests_count, extra_guests_count - ) + check_spare_licenses_available(realm, plan, extra_non_guests_count, extra_guests_count) except LicenseLimitError: if num_invites == 1: message = _("All Zulip licenses for this organization are currently in use.")