ref: Remove PlanData references for Plan model properties (#1150)

ajay-sentry · web-flow · commit 06167fb1fd01 · 2025-02-11T23:52:54.000Z
diff --git a/api/internal/owner/serializers.py b/api/internal/owner/serializers.py
@@ -131,7 +131,7 @@ def validate_value(self, value: str) -> str:
 
         plan_service = PlanService(current_org=current_org)
         plan_values = [
-            plan["value"] for plan in plan_service.available_plans(current_owner)
+            plan.name for plan in plan_service.available_plans(current_owner)
         ]
         if value not in plan_values:
             raise serializers.ValidationError(
diff --git a/api/internal/tests/views/test_user_viewset.py b/api/internal/tests/views/test_user_viewset.py
@@ -10,7 +10,7 @@
     PullFactory,
     RepositoryFactory,
 )
-from shared.plan.constants import PlanName, TierName
+from shared.plan.constants import DEFAULT_FREE_PLAN, TierName
 
 from core.models import Pull
 from utils.test_utils import APIClient
@@ -20,7 +20,7 @@ class UserViewSetTests(APITestCase):
     def setUp(self):
         non_org_active_user = OwnerFactory()
         tier = TierFactory(tier_name=TierName.BASIC.value)
-        plan = PlanFactory(name=PlanName.BASIC_PLAN_NAME.value, tier=tier)
+        plan = PlanFactory(name=DEFAULT_FREE_PLAN, tier=tier)
         self.current_owner = OwnerFactory(
             plan=plan.name,
             plan_user_count=5,
diff --git a/billing/tests/test_views.py b/billing/tests/test_views.py
@@ -871,7 +871,7 @@ def test_customer_subscription_updated_does_not_change_subscription_if_not_paid_
         has_unverified_initial_payment_method_mock,
     ):
         has_unverified_initial_payment_method_mock.return_value = False
-        self.owner.plan = PlanName.BASIC_PLAN_NAME.value
+        self.owner.plan = DEFAULT_FREE_PLAN
         self.owner.plan_user_count = 0
         self.owner.plan_auto_activate = False
         self.owner.save()
@@ -895,7 +895,7 @@ def test_customer_subscription_updated_does_not_change_subscription_if_not_paid_
         )
 
         self.owner.refresh_from_db()
-        assert self.owner.plan == PlanName.BASIC_PLAN_NAME.value
+        assert self.owner.plan == DEFAULT_FREE_PLAN
         assert self.owner.plan_user_count == 0
         assert self.owner.plan_auto_activate == False
         pm_mock.assert_called_once_with(
diff --git a/codecov_auth/tests/unit/services/test_org_level_token_service.py b/codecov_auth/tests/unit/services/test_org_level_token_service.py
@@ -10,7 +10,7 @@
     PlanFactory,
     TierFactory,
 )
-from shared.plan.constants import PlanName, TierName
+from shared.plan.constants import DEFAULT_FREE_PLAN, PlanName, TierName
 
 from codecov_auth.models import OrganizationLevelToken
 from codecov_auth.services.org_level_token_service import OrgLevelTokenService
@@ -47,7 +47,7 @@ def setUp(self):
         self.basic_tier = TierFactory(tier_name=TierName.BASIC.value)
         self.basic_plan = PlanFactory(
             tier=self.basic_tier,
-            name=PlanName.BASIC_PLAN_NAME.value,
+            name=DEFAULT_FREE_PLAN,
         )
         self.owner = OwnerFactory(plan=self.enterprise_plan.name)
 
diff --git a/graphql_api/tests/test_coverage_analytics.py b/graphql_api/tests/test_coverage_analytics.py
@@ -10,6 +10,7 @@
     RepositoryFactory,
 )
 
+from billing.helpers import mock_all_plans_and_tiers
 from core.models import Commit, Repository
 from graphql_api.tests.helper import GraphQLTestHelper
 from graphql_api.types.coverage_analytics.coverage_analytics import (
@@ -248,6 +249,7 @@ def test_coverage_analytics_resolves_to_error(self) -> None:
     def test_coverage_analytics_with_interval(self):
         """Test with interval argument to fetch coverage data in a specific time range"""
 
+        mock_all_plans_and_tiers()
         # Create data to populate the timeseries graph
         repo = self.create_repository("test-repo")
         one_day_ago = timezone.make_aware(datetime.datetime(2022, 1, 1, 0, 0))
diff --git a/graphql_api/tests/test_owner.py b/graphql_api/tests/test_owner.py
@@ -1198,7 +1198,7 @@ def test_fetch_owner_with_no_service(self):
         current_org = OwnerFactory(
             username="random-plan-user",
             service="github",
-            plan=PlanName.BASIC_PLAN_NAME.value,
+            plan=DEFAULT_FREE_PLAN,
         )
 
         query = """{
diff --git a/graphql_api/types/owner/owner.py b/graphql_api/types/owner/owner.py
@@ -8,7 +8,7 @@
 from ariadne import ObjectType
 from django.conf import settings
 from graphql import GraphQLResolveInfo
-from shared.plan.constants import DEFAULT_FREE_PLAN, PlanData, convert_to_DTO
+from shared.plan.constants import DEFAULT_FREE_PLAN
 from shared.plan.service import PlanService
 
 import services.activation as activation
@@ -111,16 +111,16 @@ def resolve_plan(owner: Owner, info: GraphQLResolveInfo) -> PlanService:
 @owner_bindable.field("pretrialPlan")
 @require_part_of_org
 @sync_to_async
-def resolve_plan_representation(owner: Owner, info: GraphQLResolveInfo) -> PlanData:
+def resolve_plan_representation(owner: Owner, info: GraphQLResolveInfo) -> Plan:
     info.context["plan_service"] = PlanService(current_org=owner)
     free_plan = Plan.objects.select_related("tier").get(name=DEFAULT_FREE_PLAN)
-    return convert_to_DTO(free_plan)
+    return free_plan
 
 
 @owner_bindable.field("availablePlans")
 @require_part_of_org
 @sync_to_async
-def resolve_available_plans(owner: Owner, info: GraphQLResolveInfo) -> List[PlanData]:
+def resolve_available_plans(owner: Owner, info: GraphQLResolveInfo) -> List[Plan]:
     plan_service = PlanService(current_org=owner)
     info.context["plan_service"] = plan_service
     owner = info.context["request"].current_owner
@@ -130,7 +130,7 @@ def resolve_available_plans(owner: Owner, info: GraphQLResolveInfo) -> List[Plan
 @owner_bindable.field("hasPrivateRepos")
 @sync_to_async
 @require_part_of_org
-def resolve_has_private_repos(owner: Owner, info: GraphQLResolveInfo) -> List[PlanData]:
+def resolve_has_private_repos(owner: Owner, info: GraphQLResolveInfo) -> bool:
     return owner.has_private_repos
 
 
diff --git a/graphql_api/types/plan_representation/plan_representation.py b/graphql_api/types/plan_representation/plan_representation.py
@@ -1,10 +1,10 @@
 from typing import List, Optional
 
 from ariadne import ObjectType
-from shared.plan.constants import PlanData
 from shared.plan.service import PlanService
 
 from codecov.db import sync_to_async
+from codecov_auth.models import Plan
 from graphql_api.helpers.ariadne import ariadne_load_local_graphql
 
 plan_representation = ariadne_load_local_graphql(
@@ -14,70 +14,70 @@
 
 
 @plan_representation_bindable.field("marketingName")
-def resolve_marketing_name(plan_data: PlanData, info) -> str:
-    return plan_data["marketing_name"]
+def resolve_marketing_name(plan_data: Plan, info) -> str:
+    return plan_data.marketing_name
 
 
 @plan_representation_bindable.field("value")
-def resolve_plan_value(plan_data: PlanData, info) -> str:
-    return plan_data["value"]
+def resolve_plan_value(plan_data: Plan, info) -> str:
+    return plan_data.name
 
 
 @plan_representation_bindable.field("billingRate")
-def resolve_billing_rate(plan_data: PlanData, info) -> Optional[str]:
-    return plan_data["billing_rate"]
+def resolve_billing_rate(plan_data: Plan, info) -> Optional[str]:
+    return plan_data.billing_rate
 
 
 @plan_representation_bindable.field("baseUnitPrice")
-def resolve_base_unit_price(plan_data: PlanData, info) -> int:
-    return plan_data["base_unit_price"]
+def resolve_base_unit_price(plan_data: Plan, info) -> int:
+    return plan_data.base_unit_price
 
 
 @plan_representation_bindable.field("benefits")
 @sync_to_async
-def resolve_benefits(plan_data: PlanData, info) -> List[str]:
+def resolve_benefits(plan_data: Plan, info) -> List[str]:
     plan_service: PlanService = info.context["plan_service"]
     if plan_service.is_org_trialing:
         benefits_with_pretrial_users = [
             benefit.replace(
                 "Up to 1 user", f"Up to {plan_service.pretrial_users_count} users"
             )
-            for benefit in plan_data["benefits"]
+            for benefit in plan_data.benefits
         ]
         return benefits_with_pretrial_users
-    return plan_data["benefits"]
+    return plan_data.benefits
 
 
 @plan_representation_bindable.field("monthlyUploadLimit")
-def resolve_monthly_uploads_limit(plan_data: PlanData, info) -> Optional[int]:
-    return plan_data["monthly_uploads_limit"]
+def resolve_monthly_uploads_limit(plan_data: Plan, info) -> Optional[int]:
+    return plan_data.monthly_uploads_limit
 
 
 @plan_representation_bindable.field("isEnterprisePlan")
-def resolve_is_enterprise(plan_data: PlanData, info) -> bool:
-    return plan_data["is_enterprise_plan"]
+def resolve_is_enterprise(plan_data: Plan, info) -> bool:
+    return plan_data.is_enterprise_plan
 
 
 @plan_representation_bindable.field("isFreePlan")
-def resolve_is_free(plan_data: PlanData, info) -> bool:
-    return plan_data["is_free_plan"]
+def resolve_is_free(plan_data: Plan, info) -> bool:
+    return plan_data.is_free_plan
 
 
 @plan_representation_bindable.field("isProPlan")
-def resolve_is_pro(plan_data: PlanData, info) -> bool:
-    return plan_data["is_pro_plan"]
+def resolve_is_pro(plan_data: Plan, info) -> bool:
+    return plan_data.is_pro_plan
 
 
 @plan_representation_bindable.field("isTeamPlan")
-def resolve_is_team(plan_data: PlanData, info) -> bool:
-    return plan_data["is_team_plan"]
+def resolve_is_team(plan_data: Plan, info) -> bool:
+    return plan_data.is_team_plan
 
 
 @plan_representation_bindable.field("isSentryPlan")
-def resolve_is_sentry(plan_data: PlanData, info) -> bool:
-    return plan_data["is_sentry_plan"]
+def resolve_is_sentry(plan_data: Plan, info) -> bool:
+    return plan_data.is_sentry_plan
 
 
 @plan_representation_bindable.field("isTrialPlan")
-def resolve_is_trial(plan_data: PlanData, info) -> bool:
-    return plan_data["is_trial_plan"]
+def resolve_is_trial(plan_data: Plan, info) -> bool:
+    return plan_data.is_trial_plan
diff --git a/requirements.in b/requirements.in
@@ -25,7 +25,7 @@ freezegun
 google-cloud-pubsub
 gunicorn>=22.0.0
 https://github.com/codecov/opentelem-python/archive/refs/tags/v0.0.4a1.tar.gz#egg=codecovopentelem
-https://github.com/codecov/shared/archive/69b736859ab71e1e3e91018439a6cb6b792502f3.tar.gz#egg=shared
+https://github.com/codecov/shared/archive/016a756f2ab982016bc2d46e2467be71604c0ba5.tar.gz#egg=shared
 https://github.com/photocrowd/django-cursor-pagination/archive/f560902696b0c8509e4d95c10ba0d62700181d84.tar.gz
 idna>=3.7
 minio
diff --git a/requirements.txt b/requirements.txt
@@ -416,7 +416,7 @@ sentry-sdk[celery]==2.13.0
     #   shared
 setproctitle==1.1.10
     # via -r requirements.in
-shared @ https://github.com/codecov/shared/archive/69b736859ab71e1e3e91018439a6cb6b792502f3.tar.gz
+shared @ https://github.com/codecov/shared/archive/016a756f2ab982016bc2d46e2467be71604c0ba5.tar.gz
     # via -r requirements.in
 simplejson==3.17.2
     # via -r requirements.in
diff --git a/services/task/task_router.py b/services/task/task_router.py
@@ -1,6 +1,6 @@
 import shared.celery_config as shared_celery_config
-from shared.billing import BillingPlan
 from shared.celery_router import route_tasks_based_on_user_plan
+from shared.plan.constants import DEFAULT_FREE_PLAN
 
 from codecov_auth.models import Owner
 from compare.models import CommitComparison
@@ -14,14 +14,14 @@ def _get_user_plan_from_ownerid(ownerid, *args, **kwargs) -> str:
     owner = Owner.objects.filter(ownerid=ownerid).first()
     if owner:
         return owner.plan
-    return BillingPlan.users_basic.db_name
+    return DEFAULT_FREE_PLAN
 
 
 def _get_user_plan_from_repoid(repoid, *args, **kwargs) -> str:
     repo = Repository.objects.filter(repoid=repoid).first()
     if repo and repo.author:
         return repo.author.plan
-    return BillingPlan.users_basic.db_name
+    return DEFAULT_FREE_PLAN
 
 
 def _get_user_plan_from_profiling_commit(profiling_id, *args, **kwargs) -> str:
@@ -32,7 +32,7 @@ def _get_user_plan_from_profiling_commit(profiling_id, *args, **kwargs) -> str:
         and profiling_commit.repository.author
     ):
         return profiling_commit.repository.author.plan
-    return BillingPlan.users_basic.db_name
+    return DEFAULT_FREE_PLAN
 
 
 def _get_user_plan_from_profiling_upload(profiling_upload_id, *args, **kwargs) -> str:
@@ -44,7 +44,7 @@ def _get_user_plan_from_profiling_upload(profiling_upload_id, *args, **kwargs) -
         and profiling_upload.profiling_commit.repository.author
     ):
         return profiling_upload.profiling_commit.repository.author.plan
-    return BillingPlan.users_basic.db_name
+    return DEFAULT_FREE_PLAN
 
 
 def _get_user_plan_from_comparison_id(comparison_id, *args, **kwargs) -> str:
@@ -60,7 +60,7 @@ def _get_user_plan_from_comparison_id(comparison_id, *args, **kwargs) -> str:
         and compare_commit.compare_commit.repository.author
     ):
         return compare_commit.compare_commit.repository.author.plan
-    return BillingPlan.users_basic.db_name
+    return DEFAULT_FREE_PLAN
 
 
 def _get_user_plan_from_label_request_id(request_id, *args, **kwargs) -> str:
@@ -76,7 +76,7 @@ def _get_user_plan_from_label_request_id(request_id, *args, **kwargs) -> str:
         and label_analysis_request.head_commit.repository.author
     ):
         return label_analysis_request.head_commit.repository.author.plan
-    return BillingPlan.users_basic.db_name
+    return DEFAULT_FREE_PLAN
 
 
 def _get_user_plan_from_suite_id(suite_id, *args, **kwargs) -> str:
@@ -92,7 +92,7 @@ def _get_user_plan_from_suite_id(suite_id, *args, **kwargs) -> str:
         and static_analysis_suite.commit.repository.author
     ):
         return static_analysis_suite.commit.repository.author.plan
-    return BillingPlan.users_basic.db_name
+    return DEFAULT_FREE_PLAN
 
 
 def _get_user_plan_from_task(task_name: str, task_kwargs: dict) -> str:
@@ -119,7 +119,7 @@ def _get_user_plan_from_task(task_name: str, task_kwargs: dict) -> str:
         shared_celery_config.static_analysis_task_name: _get_user_plan_from_suite_id,
     }
     func_to_use = owner_plan_lookup_funcs.get(
-        task_name, lambda *args, **kwargs: BillingPlan.users_basic.db_name
+        task_name, lambda *args, **kwargs: DEFAULT_FREE_PLAN
     )
     return func_to_use(**task_kwargs)
 
diff --git a/services/tests/test_task_router.py b/services/tests/test_task_router.py
diff --git a/upload/tests/test_serializers.py b/upload/tests/test_serializers.py
diff --git a/upload/tests/views/test_upload_coverage.py b/upload/tests/views/test_upload_coverage.py
diff --git a/webhook_handlers/tests/test_github.py b/webhook_handlers/tests/test_github.py
diff --git a/webhook_handlers/tests/test_github_enterprise.py b/webhook_handlers/tests/test_github_enterprise.py

Original file line number	Diff line number	Diff line change
`@@ -131,7 +131,7 @@ def validate_value(self, value: str) -> str:`
`131`	`131`
`132`	`132`	`plan_service = PlanService(current_org=current_org)`
`133`	`133`	`plan_values = [`
`134`		`- plan["value"] for plan in plan_service.available_plans(current_owner)`
	`134`	`+ plan.name for plan in plan_service.available_plans(current_owner)`
`135`	`135`	`]`
`136`	`136`	`if value not in plan_values:`
`137`	`137`	`raise serializers.ValidationError(`
Original file line number	Diff line number	Diff line change
`@@ -1198,7 +1198,7 @@ def test_fetch_owner_with_no_service(self):`
`1198`	`1198`	`current_org = OwnerFactory(`
`1199`	`1199`	`username="random-plan-user",`
`1200`	`1200`	`service="github",`
`1201`		`- plan=PlanName.BASIC_PLAN_NAME.value,`
	`1201`	`+ plan=DEFAULT_FREE_PLAN,`
`1202`	`1202`	`)`
`1203`	`1203`
`1204`	`1204`	`query = """{`