Skip to content
This repository was archived by the owner on May 5, 2025. It is now read-only.

Commit 016a756

Browse files
ajay-sentryajay-sentry
authored
remove billing plan enum and update tests (#510)
1 parent 33b9af9 commit 016a756

File tree

6 files changed

+77
-207
lines changed

6 files changed

+77
-207
lines changed

‎shared/billing/__init__.py‎

Lines changed: 0 additions & 53 deletions
This file was deleted.

‎shared/celery_router.py‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -3,9 +3,9 @@
33
from collections import OrderedDict
44
from collections.abc import Mapping
55

6-
from shared.billing import BillingPlan, is_enterprise_cloud_plan
76
from shared.celery_config import BaseCeleryConfig, get_task_group
87
from shared.config import get_config
8+
from shared.django_apps.codecov_auth.models import Plan
99

1010
Pattern = re.Pattern
1111

@@ -48,8 +48,8 @@ def route_tasks_based_on_user_plan(task_name: str, user_plan: str):
4848
default_task_queue = (
4949
route(task_name) or dict(queue=BaseCeleryConfig.task_default_queue)
5050
)["queue"]
51-
billing_plan = BillingPlan.from_str(user_plan)
52-
if is_enterprise_cloud_plan(billing_plan):
51+
plan = Plan.objects.get(name=user_plan)
52+
if plan.is_enterprise_plan:
5353
default_enterprise_queue_specific_config = get_config(
5454
"setup", "tasks", "celery", "enterprise", default=dict()
5555
)

‎shared/plan/service.py‎

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3,10 +3,12 @@
33
from functools import cached_property
44
from typing import List, Optional
55

6-
from shared.billing import is_pr_billing_plan
6+
from django.conf import settings
7+
78
from shared.config import get_config
89
from shared.django_apps.codecov.commands.exceptions import ValidationError
910
from shared.django_apps.codecov_auth.models import Owner, Plan, Service
11+
from shared.license import get_current_license
1012
from shared.plan.constants import (
1113
DEFAULT_FREE_PLAN,
1214
TEAM_PLAN_MAX_USERS,
@@ -348,4 +350,10 @@ def is_trial_plan(self) -> bool:
348350

349351
@property
350352
def is_pr_billing_plan(self) -> bool:
351-
return is_pr_billing_plan(plan=self.plan_name)
353+
if not settings.IS_ENTERPRISE:
354+
return self.plan_data.name not in [
355+
PlanName.CODECOV_PRO_MONTHLY_LEGACY.value,
356+
PlanName.CODECOV_PRO_YEARLY_LEGACY.value,
357+
]
358+
else:
359+
return get_current_license().is_pr_billing

‎tests/unit/billing/__init__.py‎

Whitespace-only changes.

‎tests/unit/billing/test_enum_definitions.py‎

Lines changed: 0 additions & 99 deletions
This file was deleted.

‎tests/unit/test_router.py‎

Lines changed: 64 additions & 50 deletions
Original file line numberDiff line numberDiff line change
@@ -1,59 +1,73 @@
11
from unittest.mock import patch
22

3-
from shared.billing import BillingPlan
3+
import pytest
4+
45
from shared.celery_config import timeseries_backfill_task_name, upload_task_name
56
from shared.celery_router import route_tasks_based_on_user_plan
7+
from shared.plan.constants import DEFAULT_FREE_PLAN, PlanName
8+
from tests.helper import mock_all_plans_and_tiers
69

710

8-
def test_route_tasks_based_on_user_plan_defaults():
9-
assert route_tasks_based_on_user_plan(
10-
upload_task_name, BillingPlan.users_basic.db_name
11-
) == {"queue": "celery", "extra_config": {}}
12-
assert route_tasks_based_on_user_plan(
13-
upload_task_name, BillingPlan.enterprise_cloud_monthly.db_name
14-
) == {"queue": "enterprise_celery", "extra_config": {}}
15-
assert route_tasks_based_on_user_plan(
16-
"misterious_task", BillingPlan.users_basic.db_name
17-
) == {"queue": "celery", "extra_config": {}}
18-
assert route_tasks_based_on_user_plan(
19-
"misterious_task", BillingPlan.enterprise_cloud_monthly.db_name
20-
) == {"queue": "enterprise_celery", "extra_config": {}}
21-
22-
23-
def test_route_tasks_with_config(mock_configuration):
24-
mock_configuration._params["setup"] = {
25-
"tasks": {
26-
"celery": {"enterprise": {"soft_timelimit": 100, "hard_timelimit": 200}},
27-
"timeseries": {
28-
"enterprise": {"soft_timelimit": 400, "hard_timelimit": 500}
29-
},
11+
class TestCeleryRouter:
12+
@pytest.mark.django_db
13+
def test_route_tasks_based_on_user_plan_defaults(self):
14+
mock_all_plans_and_tiers()
15+
assert route_tasks_based_on_user_plan(upload_task_name, DEFAULT_FREE_PLAN) == {
16+
"queue": "celery",
17+
"extra_config": {},
3018
}
31-
}
32-
assert route_tasks_based_on_user_plan(
33-
upload_task_name, BillingPlan.users_basic.db_name
34-
) == {"queue": "celery", "extra_config": {}}
35-
assert route_tasks_based_on_user_plan(
36-
upload_task_name, BillingPlan.enterprise_cloud_monthly.db_name
37-
) == {
38-
"queue": "enterprise_celery",
39-
"extra_config": {"soft_timelimit": 100, "hard_timelimit": 200},
40-
}
41-
assert route_tasks_based_on_user_plan(
42-
timeseries_backfill_task_name, BillingPlan.enterprise_cloud_monthly.db_name
43-
) == {
44-
"queue": "enterprise_celery",
45-
"extra_config": {"soft_timelimit": 400, "hard_timelimit": 500},
46-
}
19+
assert route_tasks_based_on_user_plan(
20+
upload_task_name, PlanName.ENTERPRISE_CLOUD_MONTHLY.value
21+
) == {"queue": "enterprise_celery", "extra_config": {}}
22+
assert route_tasks_based_on_user_plan("misterious_task", DEFAULT_FREE_PLAN) == {
23+
"queue": "celery",
24+
"extra_config": {},
25+
}
26+
assert route_tasks_based_on_user_plan(
27+
"misterious_task", PlanName.ENTERPRISE_CLOUD_MONTHLY.value
28+
) == {"queue": "enterprise_celery", "extra_config": {}}
4729

30+
@pytest.mark.django_db
31+
def test_route_tasks_with_config(self, mock_configuration):
32+
mock_all_plans_and_tiers()
33+
mock_configuration._params["setup"] = {
34+
"tasks": {
35+
"celery": {
36+
"enterprise": {"soft_timelimit": 100, "hard_timelimit": 200}
37+
},
38+
"timeseries": {
39+
"enterprise": {"soft_timelimit": 400, "hard_timelimit": 500}
40+
},
41+
}
42+
}
43+
assert route_tasks_based_on_user_plan(upload_task_name, DEFAULT_FREE_PLAN) == {
44+
"queue": "celery",
45+
"extra_config": {},
46+
}
47+
assert route_tasks_based_on_user_plan(
48+
upload_task_name, PlanName.ENTERPRISE_CLOUD_MONTHLY.value
49+
) == {
50+
"queue": "enterprise_celery",
51+
"extra_config": {"soft_timelimit": 100, "hard_timelimit": 200},
52+
}
53+
assert route_tasks_based_on_user_plan(
54+
timeseries_backfill_task_name, PlanName.ENTERPRISE_CLOUD_MONTHLY.value
55+
) == {
56+
"queue": "enterprise_celery",
57+
"extra_config": {"soft_timelimit": 400, "hard_timelimit": 500},
58+
}
4859

49-
@patch.dict(
50-
"shared.celery_config.BaseCeleryConfig.task_routes",
51-
{"app.tasks.upload.*": {"queue": "uploads"}},
52-
)
53-
def test_route_tasks_with_glob_config(mocker):
54-
assert route_tasks_based_on_user_plan(
55-
upload_task_name, BillingPlan.users_basic.db_name
56-
) == {"queue": "uploads", "extra_config": {}}
57-
assert route_tasks_based_on_user_plan(
58-
upload_task_name, BillingPlan.enterprise_cloud_monthly.db_name
59-
) == {"queue": "enterprise_uploads", "extra_config": {}}
60+
@patch.dict(
61+
"shared.celery_config.BaseCeleryConfig.task_routes",
62+
{"app.tasks.upload.*": {"queue": "uploads"}},
63+
)
64+
@pytest.mark.django_db
65+
def test_route_tasks_with_glob_config(self, mocker):
66+
mock_all_plans_and_tiers()
67+
assert route_tasks_based_on_user_plan(upload_task_name, DEFAULT_FREE_PLAN) == {
68+
"queue": "uploads",
69+
"extra_config": {},
70+
}
71+
assert route_tasks_based_on_user_plan(
72+
upload_task_name, PlanName.ENTERPRISE_CLOUD_MONTHLY.value
73+
) == {"queue": "enterprise_uploads", "extra_config": {}}

0 commit comments

Comments
 (0)