realm-settings: Get max invites for realm plan type helper.

There are a few places where we want to set the max invites for a
realm to the default for a realm's plan type, so this creates a
helper function that can be used consistently to get that default
value.

corporate/views/support.py

@@ -75,7 +75,11 @@ from zerver.models import (
     RealmReactivationStatus,
     UserProfile,
 )
-from zerver.models.realms import get_org_type_display_name, get_realm
+from zerver.models.realms import (
+    get_default_max_invites_for_realm_plan_type,
+    get_org_type_display_name,
+    get_realm,
+)
 from zerver.models.users import get_user_profile_by_id
 from zerver.views.invite import get_invitee_emails_set
 from zilencer.lib.remote_counts import MissingDataError, compute_max_monthly_messages
@@ -314,13 +318,10 @@ def get_realm_plan_type_options_for_discount() -> list[SupportSelectOption]:
 
 
 def get_default_max_invites_for_plan_type(realm: Realm) -> int:
-    if realm.plan_type in [
+    default_max = get_default_max_invites_for_realm_plan_type(realm.plan_type)
-        Realm.PLAN_TYPE_PLUS,
+    if default_max is None:
-        Realm.PLAN_TYPE_STANDARD,
-        Realm.PLAN_TYPE_STANDARD_FREE,
-    ]:
-        return Realm.INVITES_STANDARD_REALM_DAILY_MAX
         return settings.INVITES_DEFAULT_REALM_DAILY_MAX
+    return default_max
 
 
 def check_update_max_invites(realm: Realm, new_max: int, default_max: int) -> bool:

zerver/actions/realm_settings.py

@@ -46,7 +46,7 @@ from zerver.models import (
     UserProfile,
 )
 from zerver.models.groups import SystemGroups
-from zerver.models.realms import get_realm
+from zerver.models.realms import get_default_max_invites_for_realm_plan_type, get_realm
 from zerver.models.users import active_user_ids
 from zerver.tornado.django_api import send_event_on_commit
 
@@ -700,22 +700,12 @@ def do_change_realm_org_type(
 @transaction.atomic(durable=True)
 def do_change_realm_max_invites(realm: Realm, max_invites: int, acting_user: UserProfile) -> None:
     old_value = realm.max_invites
-    new_max = max_invites
+    if max_invites == 0:
-
         # Reset to default maximum for plan type
-    if new_max == 0:
+        new_max = get_default_max_invites_for_realm_plan_type(realm.plan_type)
-        if realm.plan_type == Realm.PLAN_TYPE_PLUS:
+    else:
-            new_max = Realm.INVITES_STANDARD_REALM_DAILY_MAX
+        new_max = max_invites
-        elif realm.plan_type == Realm.PLAN_TYPE_STANDARD:
+    realm.max_invites = new_max  # type: ignore[assignment] # https://github.com/python/mypy/issues/3004
-            new_max = Realm.INVITES_STANDARD_REALM_DAILY_MAX
-        elif realm.plan_type == Realm.PLAN_TYPE_SELF_HOSTED:
-            new_max = None  # type: ignore[assignment] # https://github.com/python/mypy/issues/3004
-        elif realm.plan_type == Realm.PLAN_TYPE_STANDARD_FREE:
-            new_max = Realm.INVITES_STANDARD_REALM_DAILY_MAX
-        elif realm.plan_type == Realm.PLAN_TYPE_LIMITED:
-            new_max = settings.INVITES_DEFAULT_REALM_DAILY_MAX
-
-    realm.max_invites = new_max
     realm.save(update_fields=["_max_invites"])
 
     RealmAuditLog.objects.create(
@@ -738,6 +728,8 @@ def do_change_realm_plan_type(
     from zproject.backends import AUTH_BACKEND_NAME_MAP
 
     old_value = realm.plan_type
+    if plan_type not in Realm.ALL_PLAN_TYPES:
+        raise AssertionError("Invalid plan type")
 
     if plan_type == Realm.PLAN_TYPE_LIMITED:
         # We do not allow public access on limited plans.
@@ -781,23 +773,11 @@ def do_change_realm_plan_type(
         extra_data={"old_value": old_value, "new_value": plan_type},
     )
 
-    if plan_type == Realm.PLAN_TYPE_PLUS:
+    realm.max_invites = get_default_max_invites_for_realm_plan_type(plan_type)  # type: ignore[assignment] # https://github.com/python/mypy/issues/3004
-        realm.max_invites = Realm.INVITES_STANDARD_REALM_DAILY_MAX
+    if plan_type == Realm.PLAN_TYPE_LIMITED:
-        realm.message_visibility_limit = None
-    elif plan_type == Realm.PLAN_TYPE_STANDARD:
-        realm.max_invites = Realm.INVITES_STANDARD_REALM_DAILY_MAX
-        realm.message_visibility_limit = None
-    elif plan_type == Realm.PLAN_TYPE_SELF_HOSTED:
-        realm.max_invites = None  # type: ignore[assignment] # https://github.com/python/mypy/issues/3004
-        realm.message_visibility_limit = None
-    elif plan_type == Realm.PLAN_TYPE_STANDARD_FREE:
-        realm.max_invites = Realm.INVITES_STANDARD_REALM_DAILY_MAX
-        realm.message_visibility_limit = None
-    elif plan_type == Realm.PLAN_TYPE_LIMITED:
-        realm.max_invites = settings.INVITES_DEFAULT_REALM_DAILY_MAX
         realm.message_visibility_limit = Realm.MESSAGE_VISIBILITY_LIMITED
     else:
-        raise AssertionError("Invalid plan type")
+        realm.message_visibility_limit = None
 
     update_first_visible_message_id(realm)
 

zerver/models/realms.py

@@ -555,7 +555,7 @@ class Realm(models.Model):  # type: ignore[django-manager-missing] # django-stub
     PLAN_TYPE_STANDARD_FREE = 4
     PLAN_TYPE_PLUS = 10
 
-    # Used for creating realms with different plan types.
+    # Used to check valid plan_type values and when populating test billing realms.
     ALL_PLAN_TYPES = {
         PLAN_TYPE_SELF_HOSTED: "self-hosted-plan",
         PLAN_TYPE_LIMITED: "limited-plan",
@@ -1210,6 +1210,19 @@ def get_corresponding_policy_value_for_group_setting(
     return Realm.POLICY_MEMBERS_ONLY
 
 
+def get_default_max_invites_for_realm_plan_type(plan_type: int) -> int | None:
+    assert plan_type in Realm.ALL_PLAN_TYPES
+    if plan_type in [
+        Realm.PLAN_TYPE_PLUS,
+        Realm.PLAN_TYPE_STANDARD,
+        Realm.PLAN_TYPE_STANDARD_FREE,
+    ]:
+        return Realm.INVITES_STANDARD_REALM_DAILY_MAX
+    if plan_type == Realm.PLAN_TYPE_SELF_HOSTED:
+        return None
+    return settings.INVITES_DEFAULT_REALM_DAILY_MAX
+
+
 class RealmDomain(models.Model):
     """For an organization with emails_restricted_to_domains enabled, the list of
     allowed domains"""