add missing imports + refactor to call 1 query

Author: RulaKhaled

api/internal/owner/serializers.py

@@ -12,7 +12,7 @@
 )
 from shared.plan.service import PlanService
 
-from codecov_auth.models import Owner
+from codecov_auth.models import Owner, Plan
 from services.billing import BillingService
 from services.sentry import send_user_webhook as send_sentry_webhook
 
@@ -147,10 +147,18 @@ def validate(self, plan: Dict[str, Any]) -> Dict[str, Any]:
                 detail="You cannot update your plan manually, for help or changes to plan, connect with [email protected]"
             )
 
+        active_plans = list(
+            Plan.objects.filter(paid_plan=True, is_active=True).values_list(
+                "name", "tier"
+            )
+        )
+        active_plan_names = {name for name, _ in active_plans}
+        team_tier_plans = {
+            name for name, tier in active_plans if tier == TierName.TEAM.value
+        }
+
         # Validate quantity here because we need access to whole plan object
-        if plan["value"] in Plan.objects.filter(
-            paid_plan=True, is_active=True
-        ).values_list("name", flat=True):
+        if plan["value"] in active_plan_names:
             if "quantity" not in plan:
                 raise serializers.ValidationError(
                     "Field 'quantity' required for updating to paid plans"
@@ -179,8 +187,7 @@ def validate(self, plan: Dict[str, Any]) -> Dict[str, Any]:
                     "Quantity or plan for paid plan must be different from the existing one"
                 )
             if (
-                plan["value"]
-                in Plan.objects.filter(tier=TierName.TEAM.value, is_active=True)
+                plan["value"] in team_tier_plans
                 and plan["quantity"] > TEAM_PLAN_MAX_USERS
             ):
                 raise serializers.ValidationError(

api/internal/tests/views/test_account_viewset.py

@@ -12,9 +12,11 @@
     AccountFactory,
     InvoiceBillingFactory,
     OwnerFactory,
+    PlanFactory,
+    TierFactory,
     UserFactory,
 )
-from shared.plan.constants import PlanName, TrialStatus
+from shared.plan.constants import PlanName, TierName, TrialStatus
 from stripe import StripeError
 
 from api.internal.tests.test_utils import GetAdminProviderAdapter
@@ -187,6 +189,13 @@ def test_retrieve_account_gets_account_fields(self):
     def test_retrieve_account_gets_account_fields_when_there_are_scheduled_details(
         self, mock_retrieve_subscription, mock_retrieve_schedule
     ):
+        tier = TierFactory(tier_name=TierName.BASIC.value)
+
+        PlanFactory(
+            name="users-basic",
+            tier=tier,
+            is_active=True,
+        )
         owner = OwnerFactory(
             admins=[self.current_owner.ownerid], stripe_subscription_id="sub_123"
         )

billing/helpers.py

@@ -2,7 +2,7 @@
 from django.db.models import QuerySet
 from shared.plan.constants import TierName
 
-from codecov_auth.models import Owner
+from codecov_auth.models import Owner, Plan
 
 
 def on_enterprise_plan(owner: Owner) -> bool:

codecov_auth/services/org_level_token_service.py

@@ -5,7 +5,7 @@
 from django.dispatch import receiver
 from django.forms import ValidationError
 
-from codecov_auth.models import OrganizationLevelToken, Owner
+from codecov_auth.models import OrganizationLevelToken, Owner, Plan
 
 log = logging.getLogger(__name__)
 

graphql_api/types/owner/owner.py

@@ -22,6 +22,7 @@
     Account,
     GithubAppInstallation,
     Owner,
+    Plan,
 )
 from codecov_auth.views.okta_cloud import OKTA_SIGNED_IN_ACCOUNTS_SESSION_KEY
 from core.models import Repository

services/billing.py

@@ -8,11 +8,12 @@
 from django.conf import settings
 from shared.plan.constants import (
     PlanBillingRate,
+    TierName,
 )
 from shared.plan.service import PlanService
 
 from billing.constants import REMOVED_INVOICE_STATUSES
-from codecov_auth.models import Owner
+from codecov_auth.models import Owner, Plan
 
 log = logging.getLogger(__name__)
 

upload/helpers.py

@@ -28,6 +28,7 @@
     SERVICE_GITHUB_ENTERPRISE,
     GithubAppInstallation,
     Owner,
+    Plan,
 )
 from core.models import Commit, Repository
 from reports.models import CommitReport, ReportSession