Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

ref: Remove BillingPlan enum and update references #510

Merged
merged 1 commit into from
Feb 11, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
53 changes: 0 additions & 53 deletions shared/billing/__init__.py

This file was deleted.

6 changes: 3 additions & 3 deletions shared/celery_router.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,9 +3,9 @@
from collections import OrderedDict
from collections.abc import Mapping

from shared.billing import BillingPlan, is_enterprise_cloud_plan
from shared.celery_config import BaseCeleryConfig, get_task_group
from shared.config import get_config
from shared.django_apps.codecov_auth.models import Plan

Pattern = re.Pattern

Expand Down Expand Up @@ -48,8 +48,8 @@ def route_tasks_based_on_user_plan(task_name: str, user_plan: str):
default_task_queue = (
route(task_name) or dict(queue=BaseCeleryConfig.task_default_queue)
)["queue"]
billing_plan = BillingPlan.from_str(user_plan)
if is_enterprise_cloud_plan(billing_plan):
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

does user_plan come from PlanService? it has the logic that gets the "correct" plan if owner has GL or Account

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It's not from plan service but rather these helpers depending on the task: https://github.com/codecov/codecov-api/blob/258b8cba496f2aeaf6a8c11ae05c8b6c2a085c8c/services/task/task_router.py#L13

which seems kinda messy but I didn't wanna unwind it more than I had to haha

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

plan = Plan.objects.get(name=user_plan)
if plan.is_enterprise_plan:
default_enterprise_queue_specific_config = get_config(
"setup", "tasks", "celery", "enterprise", default=dict()
)
Expand Down
12 changes: 10 additions & 2 deletions shared/plan/service.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,10 +3,12 @@
from functools import cached_property
from typing import List, Optional

from shared.billing import is_pr_billing_plan
from django.conf import settings

from shared.config import get_config
from shared.django_apps.codecov.commands.exceptions import ValidationError
from shared.django_apps.codecov_auth.models import Owner, Plan, Service
from shared.license import get_current_license
from shared.plan.constants import (
DEFAULT_FREE_PLAN,
TEAM_PLAN_MAX_USERS,
Expand Down Expand Up @@ -348,4 +350,10 @@ def is_trial_plan(self) -> bool:

@property
def is_pr_billing_plan(self) -> bool:
return is_pr_billing_plan(plan=self.plan_name)
if not settings.IS_ENTERPRISE:
return self.plan_data.name not in [
PlanName.CODECOV_PRO_MONTHLY_LEGACY.value,
PlanName.CODECOV_PRO_YEARLY_LEGACY.value,
]
else:
return get_current_license().is_pr_billing
Empty file removed tests/unit/billing/__init__.py
Empty file.
99 changes: 0 additions & 99 deletions tests/unit/billing/test_enum_definitions.py

This file was deleted.

114 changes: 64 additions & 50 deletions tests/unit/test_router.py
Original file line number Diff line number Diff line change
@@ -1,59 +1,73 @@
from unittest.mock import patch

from shared.billing import BillingPlan
import pytest

from shared.celery_config import timeseries_backfill_task_name, upload_task_name
from shared.celery_router import route_tasks_based_on_user_plan
from shared.plan.constants import DEFAULT_FREE_PLAN, PlanName
from tests.helper import mock_all_plans_and_tiers


def test_route_tasks_based_on_user_plan_defaults():
assert route_tasks_based_on_user_plan(
upload_task_name, BillingPlan.users_basic.db_name
) == {"queue": "celery", "extra_config": {}}
assert route_tasks_based_on_user_plan(
upload_task_name, BillingPlan.enterprise_cloud_monthly.db_name
) == {"queue": "enterprise_celery", "extra_config": {}}
assert route_tasks_based_on_user_plan(
"misterious_task", BillingPlan.users_basic.db_name
) == {"queue": "celery", "extra_config": {}}
assert route_tasks_based_on_user_plan(
"misterious_task", BillingPlan.enterprise_cloud_monthly.db_name
) == {"queue": "enterprise_celery", "extra_config": {}}


def test_route_tasks_with_config(mock_configuration):
mock_configuration._params["setup"] = {
"tasks": {
"celery": {"enterprise": {"soft_timelimit": 100, "hard_timelimit": 200}},
"timeseries": {
"enterprise": {"soft_timelimit": 400, "hard_timelimit": 500}
},
class TestCeleryRouter:
@pytest.mark.django_db
def test_route_tasks_based_on_user_plan_defaults(self):
mock_all_plans_and_tiers()
assert route_tasks_based_on_user_plan(upload_task_name, DEFAULT_FREE_PLAN) == {
"queue": "celery",
"extra_config": {},
}
}
assert route_tasks_based_on_user_plan(
upload_task_name, BillingPlan.users_basic.db_name
) == {"queue": "celery", "extra_config": {}}
assert route_tasks_based_on_user_plan(
upload_task_name, BillingPlan.enterprise_cloud_monthly.db_name
) == {
"queue": "enterprise_celery",
"extra_config": {"soft_timelimit": 100, "hard_timelimit": 200},
}
assert route_tasks_based_on_user_plan(
timeseries_backfill_task_name, BillingPlan.enterprise_cloud_monthly.db_name
) == {
"queue": "enterprise_celery",
"extra_config": {"soft_timelimit": 400, "hard_timelimit": 500},
}
assert route_tasks_based_on_user_plan(
upload_task_name, PlanName.ENTERPRISE_CLOUD_MONTHLY.value
) == {"queue": "enterprise_celery", "extra_config": {}}
assert route_tasks_based_on_user_plan("misterious_task", DEFAULT_FREE_PLAN) == {
"queue": "celery",
"extra_config": {},
}
assert route_tasks_based_on_user_plan(
"misterious_task", PlanName.ENTERPRISE_CLOUD_MONTHLY.value
) == {"queue": "enterprise_celery", "extra_config": {}}

@pytest.mark.django_db
def test_route_tasks_with_config(self, mock_configuration):
mock_all_plans_and_tiers()
mock_configuration._params["setup"] = {
"tasks": {
"celery": {
"enterprise": {"soft_timelimit": 100, "hard_timelimit": 200}
},
"timeseries": {
"enterprise": {"soft_timelimit": 400, "hard_timelimit": 500}
},
}
}
assert route_tasks_based_on_user_plan(upload_task_name, DEFAULT_FREE_PLAN) == {
"queue": "celery",
"extra_config": {},
}
assert route_tasks_based_on_user_plan(
upload_task_name, PlanName.ENTERPRISE_CLOUD_MONTHLY.value
) == {
"queue": "enterprise_celery",
"extra_config": {"soft_timelimit": 100, "hard_timelimit": 200},
}
assert route_tasks_based_on_user_plan(
timeseries_backfill_task_name, PlanName.ENTERPRISE_CLOUD_MONTHLY.value
) == {
"queue": "enterprise_celery",
"extra_config": {"soft_timelimit": 400, "hard_timelimit": 500},
}

@patch.dict(
"shared.celery_config.BaseCeleryConfig.task_routes",
{"app.tasks.upload.*": {"queue": "uploads"}},
)
def test_route_tasks_with_glob_config(mocker):
assert route_tasks_based_on_user_plan(
upload_task_name, BillingPlan.users_basic.db_name
) == {"queue": "uploads", "extra_config": {}}
assert route_tasks_based_on_user_plan(
upload_task_name, BillingPlan.enterprise_cloud_monthly.db_name
) == {"queue": "enterprise_uploads", "extra_config": {}}
@patch.dict(
"shared.celery_config.BaseCeleryConfig.task_routes",
{"app.tasks.upload.*": {"queue": "uploads"}},
)
@pytest.mark.django_db
def test_route_tasks_with_glob_config(self, mocker):
mock_all_plans_and_tiers()
assert route_tasks_based_on_user_plan(upload_task_name, DEFAULT_FREE_PLAN) == {
"queue": "uploads",
"extra_config": {},
}
assert route_tasks_based_on_user_plan(
upload_task_name, PlanName.ENTERPRISE_CLOUD_MONTHLY.value
) == {"queue": "enterprise_uploads", "extra_config": {}}