From 3a445a992b0db6c3ec4cf98bcec5ed08edf37a8e Mon Sep 17 00:00:00 2001 From: Nora Shapiro Date: Thu, 2 Jan 2025 17:08:32 -0800 Subject: [PATCH] update PlanService so GitLab orgs use the root org's plan, fold is_pr_billing_plan into PlanService --- shared/billing/__init__.py | 5 +++-- shared/plan/constants.py | 4 ++-- shared/plan/service.py | 16 +++++++++++++-- tests/unit/plan/test_plan.py | 38 +++++++++++++++++++++++++++++++++++- 4 files changed, 56 insertions(+), 7 deletions(-) diff --git a/shared/billing/__init__.py b/shared/billing/__init__.py index 4d508b0b..9076c1b3 100644 --- a/shared/billing/__init__.py +++ b/shared/billing/__init__.py @@ -7,8 +7,8 @@ class BillingPlan(Enum): users_ghm = "users" - users_monthly = "users-inappm" - users_yearly = "users-inappy" + users_monthly = "users-inappm" # not pr_billing_plan + users_yearly = "users-inappy" # not pr_billing_plan users_free = "users-free" users_basic = "users-basic" users_trial = "users-trial" @@ -37,6 +37,7 @@ def is_enterprise_cloud_plan(plan: BillingPlan) -> bool: def is_pr_billing_plan(plan: str) -> bool: + # use is_pr_billing_plan() in PlanService instead of accessing this directly if not settings.IS_ENTERPRISE: return plan in [ BillingPlan.pr_monthly.value, diff --git a/shared/plan/constants.py b/shared/plan/constants.py index 6823ef85..07b74e34 100644 --- a/shared/plan/constants.py +++ b/shared/plan/constants.py @@ -29,8 +29,8 @@ class PlanName(enum.Enum): TRIAL_PLAN_NAME = "users-trial" CODECOV_PRO_MONTHLY = "users-pr-inappm" CODECOV_PRO_YEARLY = "users-pr-inappy" - SENTRY_MONTHLY = "users-sentrym" - SENTRY_YEARLY = "users-sentryy" + SENTRY_MONTHLY = "users-sentrym" # not in BillingPlan + SENTRY_YEARLY = "users-sentryy" # not in BillingPlan TEAM_MONTHLY = "users-teamm" TEAM_YEARLY = "users-teamy" GHM_PLAN_NAME = "users" diff --git a/shared/plan/service.py b/shared/plan/service.py index 55f81eb4..cf0b8b23 100644 --- a/shared/plan/service.py +++ b/shared/plan/service.py @@ -2,9 +2,10 @@ from datetime import datetime, timedelta from typing import List, Optional +from shared.billing import is_pr_billing_plan from shared.config import get_config from shared.django_apps.codecov.commands.exceptions import ValidationError -from shared.django_apps.codecov_auth.models import Owner +from shared.django_apps.codecov_auth.models import Owner, Service from shared.plan.constants import ( BASIC_PLAN, ENTERPRISE_CLOUD_USER_PLAN_REPRESENTATIONS, @@ -46,7 +47,14 @@ def __init__(self, current_org: Owner): Raises: ValueError: If the organization's plan is unsupported. """ - self.current_org = current_org + if ( + current_org.service == Service.GITLAB.value + and current_org.parent_service_id + ): + # for GitLab groups and subgroups, use the plan on the root org + self.current_org = current_org.root_organization + else: + self.current_org = current_org if self.current_org.plan not in USER_PLAN_REPRESENTATIONS: raise ValueError("Unsupported plan") self._plan_data = None @@ -340,3 +348,7 @@ def is_team_plan(self) -> bool: @property def is_trial_plan(self) -> bool: return self.plan_name in TRIAL_PLAN_REPRESENTATION + + @property + def is_pr_billing_plan(self) -> bool: + return is_pr_billing_plan(plan=self.plan_name) diff --git a/tests/unit/plan/test_plan.py b/tests/unit/plan/test_plan.py index c82cd105..adeeaaa3 100644 --- a/tests/unit/plan/test_plan.py +++ b/tests/unit/plan/test_plan.py @@ -1,10 +1,11 @@ from datetime import datetime, timedelta from unittest.mock import patch -from django.test import TestCase +from django.test import TestCase, override_settings from freezegun import freeze_time from shared.django_apps.codecov.commands.exceptions import ValidationError +from shared.django_apps.codecov_auth.models import Service from shared.django_apps.codecov_auth.tests.factories import OwnerFactory from shared.plan.constants import ( BASIC_PLAN, @@ -317,6 +318,34 @@ def test_plan_service_returns_if_owner_has_trial_dates(self): assert plan_service.has_trial_dates == True + def test_plan_service_gitlab_with_root_org(self): + root_owner_org = OwnerFactory( + service=Service.GITLAB.value, + plan=PlanName.FREE_PLAN_NAME.value, + plan_user_count=1, + service_id="1234", + ) + middle_org = OwnerFactory( + service=Service.GITLAB.value, + service_id="5678", + parent_service_id=root_owner_org.service_id, + ) + child_owner_org = OwnerFactory( + service=Service.GITLAB.value, + plan=PlanName.CODECOV_PRO_MONTHLY.value, + plan_user_count=20, + parent_service_id=middle_org.service_id, + ) + # root_plan and child_plan should be the same + root_plan = PlanService(current_org=root_owner_org) + child_plan = PlanService(current_org=child_owner_org) + + assert root_plan.is_pro_plan == child_plan.is_pro_plan == False + assert root_plan.plan_user_count == child_plan.plan_user_count == 1 + assert ( + root_plan.plan_name == child_plan.plan_name == PlanName.FREE_PLAN_NAME.value + ) + class AvailablePlansBeforeTrial(TestCase): """ @@ -815,6 +844,7 @@ def test_sentry_user(self, is_sentry_user): assert self.plan_service.available_plans(owner=self.owner) == expected_result +@override_settings(IS_ENTERPRISE=False) class PlanServiceIs___PlanTests(TestCase): def test_is_trial_plan(self): self.current_org = OwnerFactory( @@ -834,6 +864,7 @@ def test_is_trial_plan(self): assert self.plan_service.is_free_plan == False assert self.plan_service.is_pro_plan == False assert self.plan_service.is_enterprise_plan == False + assert self.plan_service.is_pr_billing_plan == True def test_is_team_plan(self): self.current_org = OwnerFactory( @@ -849,6 +880,7 @@ def test_is_team_plan(self): assert self.plan_service.is_free_plan == False assert self.plan_service.is_pro_plan == False assert self.plan_service.is_enterprise_plan == False + assert self.plan_service.is_pr_billing_plan == True def test_is_sentry_plan(self): self.current_org = OwnerFactory( @@ -864,6 +896,7 @@ def test_is_sentry_plan(self): assert self.plan_service.is_free_plan == False assert self.plan_service.is_pro_plan == True assert self.plan_service.is_enterprise_plan == False + assert self.plan_service.is_pr_billing_plan == False def test_is_free_plan(self): self.current_org = OwnerFactory( @@ -878,6 +911,7 @@ def test_is_free_plan(self): assert self.plan_service.is_free_plan == True assert self.plan_service.is_pro_plan == False assert self.plan_service.is_enterprise_plan == False + assert self.plan_service.is_pr_billing_plan == True def test_is_pro_plan(self): self.current_org = OwnerFactory( @@ -892,6 +926,7 @@ def test_is_pro_plan(self): assert self.plan_service.is_free_plan == False assert self.plan_service.is_pro_plan == True assert self.plan_service.is_enterprise_plan == False + assert self.plan_service.is_pr_billing_plan == True def test_is_enterprise_plan(self): self.current_org = OwnerFactory( @@ -906,3 +941,4 @@ def test_is_enterprise_plan(self): assert self.plan_service.is_free_plan == False assert self.plan_service.is_pro_plan == False assert self.plan_service.is_enterprise_plan == True + assert self.plan_service.is_pr_billing_plan == True