Worker: Migrate to Plan / Tier Tables (#1007)

Co-authored-by: Rola Abuhasna <[email protected]>

Author: ajay-sentry

database/models/core.py

@@ -60,7 +60,9 @@ class Account(CodecovBaseModel, MixinBaseClassNoExternalID):
     name = Column(types.String(100), nullable=False, unique=True)
     is_active = Column(types.Boolean, nullable=False, default=True)
     plan = Column(
-        types.String(50), nullable=False, default=PlanName.BASIC_PLAN_NAME.value
+        types.String(50),
+        nullable=False,
+        default=PlanName.BASIC_PLAN_NAME.value,  # TODO: UPDATE WITH NEW FREE PLAN NAME
     )
     plan_seat_count = Column(types.SmallInteger, nullable=False, default=1)
     free_seat_count = Column(types.SmallInteger, nullable=False, default=0)

database/tests/factories/core.py

@@ -4,7 +4,8 @@
 
 import factory
 from factory import Factory
-from shared.plan.constants import PlanName
+from shared.django_apps.codecov_auth.models import Plan, Tier
+from shared.plan.constants import PlanName, TierName
 
 from database import enums, models
 from services.encryption import encryptor
@@ -300,6 +301,35 @@ class Meta:
     value = ""
 
 
+class TierFactory(Factory):
+    class Meta:
+        model = Tier
+
+    tier_name = TierName.BASIC.value
+    bundle_analysis = False
+    test_analytics = False
+    flaky_test_detection = False
+    project_coverage = False
+    private_repo_support = False
+
+
+class PlanFactory(Factory):
+    class Meta:
+        model = Plan
+
+    tier = factory.SubFactory(TierFactory)
+    base_unit_price = 0
+    benefits = factory.LazyFunction(lambda: ["Benefit 1", "Benefit 2", "Benefit 3"])
+    billing_rate = None
+    is_active = True
+    marketing_name = factory.Faker("catch_phrase")
+    max_seats = 1
+    monthly_uploads_limit = None
+    name = PlanName.BASIC_PLAN_NAME.value
+    paid_plan = False
+    stripe_id = None
+
+
 class UploadErrorFactory(Factory):
     class Meta:
         model = models.UploadError

requirements.in

@@ -1,5 +1,5 @@
 https://github.com/codecov/test-results-parser/archive/190bbc8a911099749928e13d5fe57f6027ca1e74.tar.gz#egg=test-results-parser
-https://github.com/codecov/shared/archive/ce6e1e9198b046af850606e8fe3633d9a25fea98.tar.gz#egg=shared
+https://github.com/codecov/shared/archive/74c0888070699b69a4da73f54be502ad05fde7b6.tar.gz#egg=shared
 https://github.com/codecov/timestring/archive/d37ceacc5954dff3b5bd2f887936a98a668dda42.tar.gz#egg=timestring
 asgiref>=3.7.2
 analytics-python==1.3.0b1

requirements.txt

@@ -377,7 +377,7 @@ sentry-sdk==2.13.0
     #   shared
 setuptools==75.7.0
     # via nodeenv
-shared @ https://github.com/codecov/shared/archive/ce6e1e9198b046af850606e8fe3633d9a25fea98.tar.gz#egg=shared
+shared @ https://github.com/codecov/shared/archive/74c0888070699b69a4da73f54be502ad05fde7b6.tar.gz#egg=shared
     # via -r requirements.in
 six==1.16.0
     # via

services/bundle_analysis/notify/contexts/tests/test_comment_context.py

@@ -22,6 +22,7 @@
 )
 from services.repository import EnrichedPull
 from services.seats import SeatActivationInfo, ShouldActivateSeat
+from tests.helpers import mock_all_plans_and_tiers
 
 
 class TestBundleAnalysisPRCommentNotificationContext:
@@ -272,7 +273,9 @@ def test_evaluate_should_use_upgrade_message(
         builder.evaluate_should_use_upgrade_message()
         assert builder._notification_context.should_use_upgrade_comment == expected
 
+    @pytest.mark.django_db
     def test_build_context(self, dbsession, mocker, mock_storage):
+        mock_all_plans_and_tiers()
         head_commit, base_commit = get_commit_pair(dbsession)
         repository = head_commit.repository
         head_commit_report, base_commit_report = get_report_pair(
@@ -305,7 +308,9 @@ def test_build_context(self, dbsession, mocker, mock_storage):
             == head_commit_report.external_id
         )
 
+    @pytest.mark.django_db
     def test_initialize_from_context(self, dbsession, mocker):
+        mock_all_plans_and_tiers()
         head_commit, base_commit = get_commit_pair(dbsession)
         user_yaml = UserYaml.from_dict(PATCH_CENTRIC_DEFAULT_CONFIG)
         builder = BundleAnalysisPRCommentContextBuilder().initialize(

services/bundle_analysis/notify/contexts/tests/test_commit_status_context.py

@@ -25,6 +25,7 @@
 from services.bundle_analysis.notify.types import NotificationUserConfig
 from services.repository import EnrichedPull
 from services.seats import SeatActivationInfo, ShouldActivateSeat
+from tests.helpers import mock_all_plans_and_tiers
 
 
 class TestBundleAnalysisPRCommentNotificationContext:
@@ -217,7 +218,9 @@ def test_load_commit_status_level(
         context = builder.get_result()
         assert context.commit_status_level == expected
 
+    @pytest.mark.django_db
     def test_build_context(self, dbsession, mocker, mock_storage):
+        mock_all_plans_and_tiers()
         head_commit, base_commit = get_commit_pair(dbsession)
         repository = head_commit.repository
         head_commit_report, base_commit_report = get_report_pair(
@@ -259,7 +262,9 @@ def test_build_context(self, dbsession, mocker, mock_storage):
         assert context.cache_ttl == 600
         assert context.commit_status_url is not None
 
+    @pytest.mark.django_db
     def test_initialize_from_context(self, dbsession, mocker):
+        mock_all_plans_and_tiers()
         head_commit, base_commit = get_commit_pair(dbsession)
         user_yaml = UserYaml.from_dict(PATCH_CENTRIC_DEFAULT_CONFIG)
         builder = CommitStatusNotificationContextBuilder().initialize(

services/bundle_analysis/notify/messages/tests/test_comment.py

@@ -26,10 +26,13 @@
 )
 from services.bundle_analysis.notify.types import NotificationUserConfig
 from services.notification.notifiers.base import NotificationResult
+from tests.helpers import mock_all_plans_and_tiers
 
 
 class TestCommentMesage:
+    @pytest.mark.django_db
     def test_build_message_from_samples(self, dbsession, mocker, mock_storage):
+        mock_all_plans_and_tiers()
         head_commit, base_commit = get_commit_pair(dbsession)
         repository = head_commit.repository
         head_commit_report, base_commit_report = get_report_pair(

services/bundle_analysis/notify/messages/tests/test_commit_status.py

@@ -24,6 +24,7 @@
 )
 from services.notification.notifiers.base import NotificationResult
 from services.seats import SeatActivationInfo, ShouldActivateSeat
+from tests.helpers import mock_all_plans_and_tiers
 
 
 class FakeRedis(object):
@@ -55,6 +56,10 @@ def mock_cache(mocker):
 
 
 class TestCommitStatusMessage:
+    @pytest.fixture(autouse=True)
+    def setup(self):
+        mock_all_plans_and_tiers()
+
     @pytest.mark.parametrize(
         "user_config, expected",
         [
@@ -81,6 +86,7 @@ class TestCommitStatusMessage:
             ),
         ],
     )
+    @pytest.mark.django_db
     def test_build_message_from_samples_negative_changes(
         self, user_config, expected, dbsession, mocker, mock_storage
     ):
@@ -133,6 +139,7 @@ def test_build_message_from_samples_negative_changes(
             ),
         ],
     )
+    @pytest.mark.django_db
     def test_build_message_from_samples(
         self, user_config, expected, dbsession, mocker, mock_storage
     ):
@@ -159,6 +166,7 @@ def test_build_message_from_samples(
         message = CommitStatusMessageStrategy().build_message(context)
         assert message == expected
 
+    @pytest.mark.django_db
     def _setup_send_message_tests(
         self, dbsession, mocker, torngit_ghapp_data, mock_storage
     ):
@@ -207,6 +215,7 @@ def _setup_send_message_tests(
             "Passed with Warnings - Bundle change: 95.64% (Threshold: 5.0%)",
         )
 
+    @pytest.mark.django_db
     @pytest.mark.parametrize(
         "torngit_ghapp_data",
         [
@@ -246,6 +255,7 @@ def test_send_message_success(
         # Side effect of sending message is updating the cache
         assert mock_cache.get_backend().get(strategy._cache_key(context)) == message
 
+    @pytest.mark.django_db
     def test_send_message_fail(self, dbsession, mocker, mock_storage):
         fake_repo_provider, context, message = self._setup_send_message_tests(
             dbsession, mocker, None, mock_storage
@@ -259,6 +269,7 @@ def test_send_message_fail(self, dbsession, mocker, mock_storage):
             explanation="TorngitClientError",
         )
 
+    @pytest.mark.django_db
     def test_skip_payload_unchanged(self, dbsession, mocker, mock_storage, mock_cache):
         fake_repo_provider, context, message = self._setup_send_message_tests(
             dbsession, mocker, None, mock_storage

services/bundle_analysis/notify/tests/test_notify_service.py

@@ -29,6 +29,7 @@
 )
 from services.bundle_analysis.notify.types import NotificationType
 from services.notification.notifiers.base import NotificationResult
+from tests.helpers import mock_all_plans_and_tiers
 
 
 def override_comment_builder_and_message_strategy(mocker):
@@ -114,7 +115,9 @@ def test_build_base_context(self, mocker, dbsession, mock_storage):
         assert base_context.commit_report == head_commit_report
         assert base_context.bundle_analysis_report.session_count() == 19
 
+    @pytest.mark.django_db
     def test_create_context_success(self, dbsession, mock_storage, mocker):
+        mock_all_plans_and_tiers()
         current_yaml = UserYaml.from_dict(PATCH_CENTRIC_DEFAULT_CONFIG)
         head_commit, base_commit = get_commit_pair(dbsession)
         head_commit_report, base_commit_report = get_report_pair(

services/bundle_analysis/tests/test_bundle_analysis.py

@@ -30,6 +30,7 @@
 )
 from services.repository import EnrichedPull
 from services.urls import get_bundle_analysis_pull_url
+from tests.helpers import mock_all_plans_and_tiers
 
 
 class MockBundleReport:
@@ -690,6 +691,7 @@ async def test_bundle_analysis_save_measurements_error(dbsession, mocker, mock_s
         ),
     ],
 )
+@pytest.mark.django_db
 def test_bundle_analysis_notify_bundle_summary(
     bundle_changes: list[BundleChange],
     percent_change: float,
@@ -700,6 +702,7 @@ def test_bundle_analysis_notify_bundle_summary(
     mock_storage,
     mock_repo_provider,
 ):
+    mock_all_plans_and_tiers()
     hook_mock_repo_provider(mocker, mock_repo_provider)
     base_commit = CommitFactory()
     dbsession.add(base_commit)
@@ -1483,6 +1486,7 @@ def contributing_modules(self, pr_changed_files):
         ),
     ],
 )
+@pytest.mark.django_db
 def test_bundle_analysis_notify_individual_bundle_data(
     bundle_changes: list[BundleChange],
     route_changes: Dict[str, List[RouteChange]],
@@ -1493,6 +1497,7 @@ def test_bundle_analysis_notify_individual_bundle_data(
     mock_storage,
     mock_repo_provider,
 ):
+    mock_all_plans_and_tiers()
     percent_change = 5.56
     user_config = {
         **PATCH_CENTRIC_DEFAULT_CONFIG,

services/notification/__init__.py

@@ -10,8 +10,9 @@
 
 from celery.exceptions import CeleryError, SoftTimeLimitExceeded
 from shared.config import get_config
+from shared.django_apps.codecov_auth.models import Plan
 from shared.helpers.yaml import default_if_true
-from shared.plan.constants import TEAM_PLAN_REPRESENTATIONS
+from shared.plan.constants import TierName
 from shared.torngit.base import TorngitBaseAdapter
 from shared.yaml import UserYaml
 
@@ -68,9 +69,13 @@ def __init__(
     def _should_use_status_notifier(self, status_type: StatusType) -> bool:
         owner: Owner = self.repository.owner
 
-        if owner.plan in TEAM_PLAN_REPRESENTATIONS:
-            if status_type != StatusType.PATCH.value:
-                return False
+        plan = Plan.objects.select_related("tier").get(name=owner.plan)
+
+        if (
+            plan.tier.tier_name == TierName.TEAM.value
+            and status_type != StatusType.PATCH.value
+        ):
+            return False
 
         return True
 
@@ -83,9 +88,13 @@ def _should_use_checks_notifier(self, status_type: StatusType) -> bool:
         if owner.service not in ["github", "github_enterprise"]:
             return False
 
-        if owner.plan in TEAM_PLAN_REPRESENTATIONS:
-            if status_type != StatusType.PATCH.value:
-                return False
+        plan = Plan.objects.select_related("tier").get(name=owner.plan)
+
+        if (
+            plan.tier.tier_name == TierName.TEAM.value
+            and status_type != StatusType.PATCH.value
+        ):
+            return False
 
         app_installation_filter = filter(
             lambda obj: (

services/notification/notifiers/tests/integration/test_comment.py

@@ -8,6 +8,7 @@
 from services.comparison.types import Comparison, EnrichedPull, FullCommit
 from services.decoration import Decoration
 from services.notification.notifiers.comment import CommentNotifier
+from tests.helpers import mock_all_plans_and_tiers
 
 
 @pytest.fixture
@@ -335,6 +336,11 @@ def sample_comparison_for_limited_upload(
 
 @pytest.mark.usefixtures("is_not_first_pull")
 class TestCommentNotifierIntegration(object):
+    @pytest.fixture(autouse=True)
+    def setup(self):
+        mock_all_plans_and_tiers()
+
+    @pytest.mark.django_db
     def test_notify(self, sample_comparison, codecov_vcr, mock_configuration):
         sample_comparison.context = ComparisonContext(
             all_tests_passed=True, test_results_error=None
@@ -407,6 +413,7 @@ def test_notify(self, sample_comparison, codecov_vcr, mock_configuration):
         assert result.data_sent == {"commentid": None, "message": message, "pullid": 9}
         assert result.data_received == {"id": 1699669247}
 
+    @pytest.mark.django_db
     def test_notify_test_results_error(
         self, sample_comparison, codecov_vcr, mock_configuration
     ):
@@ -482,6 +489,7 @@ def test_notify_test_results_error(
         assert result.data_sent == {"commentid": None, "message": message, "pullid": 9}
         assert result.data_received == {"id": 1699669247}
 
+    @pytest.mark.django_db
     def test_notify_upgrade(
         self, dbsession, sample_comparison_for_upgrade, codecov_vcr, mock_configuration
     ):
@@ -516,6 +524,7 @@ def test_notify_upgrade(
         }
         assert result.data_received == {"id": 1361234119}
 
+    @pytest.mark.django_db
     def test_notify_upload_limited(
         self,
         dbsession,
@@ -559,6 +568,7 @@ def test_notify_upload_limited(
         }
         assert result.data_received == {"id": 1111984446}
 
+    @pytest.mark.django_db
     def test_notify_gitlab(
         self, sample_comparison_gitlab, codecov_vcr, mock_configuration
     ):
@@ -625,6 +635,7 @@ def test_notify_gitlab(
         assert result.data_sent == {"commentid": None, "message": message, "pullid": 5}
         assert result.data_received == {"id": 1457135397}
 
+    @pytest.mark.django_db
     def test_notify_new_layout(
         self, sample_comparison, codecov_vcr, mock_configuration
     ):
@@ -695,6 +706,7 @@ def test_notify_new_layout(
         assert result.data_sent == {"commentid": None, "message": message, "pullid": 9}
         assert result.data_received == {"id": 1699669290}
 
+    @pytest.mark.django_db
     def test_notify_with_components(
         self, sample_comparison, codecov_vcr, mock_configuration
     ):