diff --git a/shared/plan/service.py b/shared/plan/service.py index 38a78b1c8..66ffdf1f3 100644 --- a/shared/plan/service.py +++ b/shared/plan/service.py @@ -7,6 +7,7 @@ from shared.django_apps.codecov_auth.models import Owner from shared.plan.constants import ( BASIC_PLAN, + ENTERPRISE_CLOUD_USER_PLAN_REPRESENTATIONS, FREE_PLAN, FREE_PLAN_REPRESENTATIONS, PR_AUTHOR_PAID_USER_PLAN_REPRESENTATIONS, @@ -16,6 +17,7 @@ TRIAL_PLAN_REPRESENTATION, TRIAL_PLAN_SEATS, USER_PLAN_REPRESENTATIONS, + PlanBillingRate, PlanData, PlanName, TrialDaysAmount, @@ -124,7 +126,7 @@ def marketing_name(self) -> str: return self.plan_data.marketing_name @property - def billing_rate(self) -> Optional[str]: + def billing_rate(self) -> Optional[PlanBillingRate]: """Returns the billing rate for the plan.""" return self.plan_data.billing_rate @@ -282,7 +284,7 @@ def trial_end_date(self) -> Optional[datetime]: return self.current_org.trial_end_date @property - def trial_total_days(self) -> Optional[int]: + def trial_total_days(self) -> Optional[TrialDaysAmount]: """Returns the total number of trial days.""" return self.plan_data.trial_days @@ -311,3 +313,30 @@ def has_seats_left(self) -> bool: self.plan_activated_users is None or len(self.plan_activated_users) < self.plan_user_count ) + + @property + def is_enterprise_plan(self) -> bool: + return self.plan_name in ENTERPRISE_CLOUD_USER_PLAN_REPRESENTATIONS + + @property + def is_free_plan(self) -> bool: + return self.plan_name in FREE_PLAN_REPRESENTATIONS + + @property + def is_pro_plan(self) -> bool: + return ( + self.plan_name in PR_AUTHOR_PAID_USER_PLAN_REPRESENTATIONS + or self.plan_name in SENTRY_PAID_USER_PLAN_REPRESENTATIONS + ) + + @property + def is_sentry_plan(self) -> bool: + return self.plan_name in SENTRY_PAID_USER_PLAN_REPRESENTATIONS + + @property + def is_team_plan(self) -> bool: + return self.plan_name in TEAM_PLAN_REPRESENTATIONS + + @property + def is_trial_plan(self) -> bool: + return self.plan_name in TRIAL_PLAN_REPRESENTATION diff --git a/tests/unit/plan/test_plan.py b/tests/unit/plan/test_plan.py index 239f10c8e..de6f48805 100644 --- a/tests/unit/plan/test_plan.py +++ b/tests/unit/plan/test_plan.py @@ -786,3 +786,96 @@ def test_sentry_user(self, is_sentry_user): # Can not do Team plan when at 11 activated users assert self.plan_service.available_plans(owner=self.owner) == expected_result + + +class PlanServiceIs___PlanTests(TestCase): + def test_is_trial_plan(self): + self.current_org = OwnerFactory( + plan=PlanName.TRIAL_PLAN_NAME.value, + trial_start_date=datetime.utcnow(), + trial_end_date=datetime.utcnow() + timedelta(days=14), + trial_status=TrialStatus.ONGOING.value, + plan_user_count=1000, + plan_activated_users=None, + ) + self.owner = OwnerFactory() + self.plan_service = PlanService(current_org=self.current_org) + + assert self.plan_service.is_trial_plan == True + assert self.plan_service.is_sentry_plan == False + assert self.plan_service.is_team_plan == False + assert self.plan_service.is_free_plan == False + assert self.plan_service.is_pro_plan == False + assert self.plan_service.is_enterprise_plan == False + + def test_is_team_plan(self): + self.current_org = OwnerFactory( + plan=PlanName.TEAM_MONTHLY.value, + trial_status=TrialStatus.EXPIRED.value, + ) + self.owner = OwnerFactory() + self.plan_service = PlanService(current_org=self.current_org) + + assert self.plan_service.is_trial_plan == False + assert self.plan_service.is_sentry_plan == False + assert self.plan_service.is_team_plan == True + assert self.plan_service.is_free_plan == False + assert self.plan_service.is_pro_plan == False + assert self.plan_service.is_enterprise_plan == False + + def test_is_sentry_plan(self): + self.current_org = OwnerFactory( + plan=PlanName.SENTRY_MONTHLY.value, + trial_status=TrialStatus.EXPIRED.value, + ) + self.owner = OwnerFactory() + self.plan_service = PlanService(current_org=self.current_org) + + assert self.plan_service.is_trial_plan == False + assert self.plan_service.is_sentry_plan == True + assert self.plan_service.is_team_plan == False + assert self.plan_service.is_free_plan == False + assert self.plan_service.is_pro_plan == True + assert self.plan_service.is_enterprise_plan == False + + def test_is_free_plan(self): + self.current_org = OwnerFactory( + plan=PlanName.FREE_PLAN_NAME.value, + ) + self.owner = OwnerFactory() + self.plan_service = PlanService(current_org=self.current_org) + + assert self.plan_service.is_trial_plan == False + assert self.plan_service.is_sentry_plan == False + assert self.plan_service.is_team_plan == False + assert self.plan_service.is_free_plan == True + assert self.plan_service.is_pro_plan == False + assert self.plan_service.is_enterprise_plan == False + + def test_is_pro_plan(self): + self.current_org = OwnerFactory( + plan=PlanName.CODECOV_PRO_MONTHLY.value, + ) + self.owner = OwnerFactory() + self.plan_service = PlanService(current_org=self.current_org) + + assert self.plan_service.is_trial_plan == False + assert self.plan_service.is_sentry_plan == False + assert self.plan_service.is_team_plan == False + assert self.plan_service.is_free_plan == False + assert self.plan_service.is_pro_plan == True + assert self.plan_service.is_enterprise_plan == False + + def test_is_enterprise_plan(self): + self.current_org = OwnerFactory( + plan=PlanName.ENTERPRISE_CLOUD_YEARLY.value, + ) + self.owner = OwnerFactory() + self.plan_service = PlanService(current_org=self.current_org) + + assert self.plan_service.is_trial_plan == False + assert self.plan_service.is_sentry_plan == False + assert self.plan_service.is_team_plan == False + assert self.plan_service.is_free_plan == False + assert self.plan_service.is_pro_plan == False + assert self.plan_service.is_enterprise_plan == True