feat: Modify billing for Team plan (#637)

api/internal/owner/serializers.py

@@ -1,14 +1,19 @@
 import logging
+from dataclasses import asdict
 from datetime import datetime
 
 from dateutil.relativedelta import relativedelta
 from django.conf import settings
 from rest_framework import serializers
 from rest_framework.exceptions import PermissionDenied
 
-from billing.helpers import available_plans
 from codecov_auth.models import Owner
-from plan.constants import PRO_PLANS, SENTRY_PAID_USER_PLAN_REPRESENTATIONS
+from plan.constants import (
+    PAID_PLANS,
+    SENTRY_PAID_USER_PLAN_REPRESENTATIONS,
+    TEAM_PLAN_MAX_USERS,
+    TEAM_PLAN_REPRESENTATIONS,
+)
 from plan.service import PlanService
 from services.billing import BillingService
 from services.sentry import send_user_webhook as send_sentry_webhook
@@ -117,8 +122,14 @@ class PlanSerializer(serializers.Serializer):
     quantity = serializers.IntegerField(required=False)
 
     def validate_value(self, value):
+        current_org = self.context["view"].owner
         current_owner = self.context["request"].current_owner
-        plan_values = [plan["value"] for plan in available_plans(current_owner)]
+
+        plan_service = PlanService(current_org=current_org)
+        available_plans = [
+            asdict(plan) for plan in plan_service.available_plans(current_owner)
+        ]
+        plan_values = [plan["value"] for plan in available_plans]
         if value not in plan_values:
             if value in SENTRY_PAID_USER_PLAN_REPRESENTATIONS:
                 log.warning(
@@ -134,7 +145,7 @@ def validate(self, plan):
         owner = self.context["view"].owner
 
         # Validate quantity here because we need access to whole plan object
-        if plan["value"] in PRO_PLANS:
+        if plan["value"] in PAID_PLANS:
             if "quantity" not in plan:
                 raise serializers.ValidationError(
                     f"Field 'quantity' required for updating to paid plans"
@@ -159,6 +170,13 @@ def validate(self, plan):
                 raise serializers.ValidationError(
                     f"Quantity or plan for paid plan must be different from the existing one"
                 )
+            if (
+                plan["value"] in TEAM_PLAN_REPRESENTATIONS
+                and plan["quantity"] > TEAM_PLAN_MAX_USERS
+            ):
+                raise serializers.ValidationError(
+                    f"Quantity for Team plan cannot exceed {TEAM_PLAN_MAX_USERS}"
+                )
         return plan
 
 
@@ -185,7 +203,7 @@ def get_plan(self, phase):
         plan_name = list(stripe_plan_dict.keys())[
             list(stripe_plan_dict.values()).index(plan_id)
         ]
-        marketing_plan_name = PRO_PLANS[plan_name].billing_rate
+        marketing_plan_name = PAID_PLANS[plan_name].billing_rate
         return marketing_plan_name
 
     def get_quantity(self, phase):

api/internal/owner/views.py

@@ -1,4 +1,5 @@
 import logging
+from dataclasses import asdict
 
 from django.db.models import F
 from django_filters import rest_framework as django_filters
@@ -10,7 +11,7 @@
 from api.shared.mixins import OwnerPropertyMixin
 from api.shared.owner.mixins import OwnerViewSetMixin, UserViewSetMixin
 from api.shared.permissions import MemberOfOrgPermissions
-from billing.helpers import available_plans, on_enterprise_plan
+from billing.helpers import on_enterprise_plan
 from services.billing import BillingService
 from services.decorators import stripe_safe
 from services.task import TaskService
@@ -151,11 +152,3 @@ def get_queryset(self):
             # pull ordering only available for enterprise
             qs = qs.annotate_last_pull_timestamp()
         return qs
-
-
-class PlanViewSet(viewsets.GenericViewSet, mixins.ListModelMixin):
-    def get_queryset(self):
-        return None
-
-    def list(self, request, *args, **kwargs):
-        return Response(available_plans(request.current_owner))

api/internal/tests/views/test_account_viewset.py

@@ -13,7 +13,7 @@
 from api.internal.tests.test_utils import GetAdminProviderAdapter
 from codecov_auth.models import Service
 from codecov_auth.tests.factories import OwnerFactory, UserFactory
-from plan.constants import PlanName
+from plan.constants import PlanName, TrialStatus
 from utils.test_utils import APIClient
 
 curr_path = os.path.dirname(__file__)
@@ -857,6 +857,111 @@ def test_update_must_fail_if_quantity_and_plan_are_equal_to_the_owners_current_o
             == "Quantity or plan for paid plan must be different from the existing one"
         )
 
+    def test_update_team_plan_must_fail_if_not_trialing(self):
+        self.current_owner.plan = PlanName.BASIC_PLAN_NAME.value
+        self.current_owner.plan_user_count = 1
+        self.current_owner.save()
+        desired_plans = [
+            {"value": PlanName.TEAM_MONTHLY.value, "quantity": 1},
+            {"value": PlanName.TEAM_YEARLY.value, "quantity": 1},
+        ]
+
+        for desired_plan in desired_plans:
+            response = self._update(
+                kwargs={
+                    "service": self.current_owner.service,
+                    "owner_username": self.current_owner.username,
+                },
+                data={"plan": desired_plan},
+            )
+
+            assert response.status_code == status.HTTP_400_BAD_REQUEST
+            assert response.json() == {
+                "plan": {
+                    "value": [
+                        f"Invalid value for plan: {desired_plan['value']}; must be one of ['users-basic', 'users-pr-inappm', 'users-pr-inappy']"
+                    ]
+                }
+            }
+
+    def test_update_team_plan_must_fail_if_too_many_activated_users_during_trial(self):
+        self.current_owner.plan = PlanName.BASIC_PLAN_NAME.value
+        self.current_owner.plan_user_count = 1
+        self.current_owner.trial_status = TrialStatus.ONGOING.value
+        self.current_owner.plan_activated_users = [i for i in range(11)]
+        self.current_owner.save()
+
+        desired_plans = [
+            {"value": PlanName.TEAM_MONTHLY.value, "quantity": 10},
+            {"value": PlanName.TEAM_YEARLY.value, "quantity": 10},
+        ]
+
+        for desired_plan in desired_plans:
+            response = self._update(
+                kwargs={
+                    "service": self.current_owner.service,
+                    "owner_username": self.current_owner.username,
+                },
+                data={"plan": desired_plan},
+            )
+
+            assert response.status_code == status.HTTP_400_BAD_REQUEST
+            assert response.json() == {
+                "plan": {
+                    "value": [
+                        f"Invalid value for plan: {desired_plan['value']}; must be one of ['users-basic', 'users-pr-inappm', 'users-pr-inappy']"
+                    ]
+                }
+            }
+
+    def test_update_team_plan_must_fail_if_currently_team_plan_add_too_many_users(self):
+        self.current_owner.plan = PlanName.TEAM_MONTHLY.value
+        self.current_owner.plan_user_count = 1
+        self.current_owner.save()
+
+        desired_plans = [
+            {"value": PlanName.TEAM_MONTHLY.value, "quantity": 11},
+            {"value": PlanName.TEAM_YEARLY.value, "quantity": 11},
+        ]
+
+        for desired_plan in desired_plans:
+            response = self._update(
+                kwargs={
+                    "service": self.current_owner.service,
+                    "owner_username": self.current_owner.username,
+                },
+                data={"plan": desired_plan},
+            )
+            assert response.status_code == status.HTTP_400_BAD_REQUEST
+            assert (
+                response.data["plan"]["non_field_errors"][0]
+                == "Quantity for Team plan cannot exceed 10"
+            )
+
+    def test_update_must_fail_if_team_plan_and_too_many_users(self):
+        desired_plans = [
+            {"value": PlanName.TEAM_MONTHLY.value, "quantity": 11},
+            {"value": PlanName.TEAM_YEARLY.value, "quantity": 11},
+        ]
+
+        for desired_plan in desired_plans:
+            response = self._update(
+                kwargs={
+                    "service": self.current_owner.service,
+                    "owner_username": self.current_owner.username,
+                },
+                data={"plan": desired_plan},
+            )
+
+            assert response.status_code == status.HTTP_400_BAD_REQUEST
+            assert response.json() == {
+                "plan": {
+                    "value": [
+                        f"Invalid value for plan: {desired_plan['value']}; must be one of ['users-basic', 'users-pr-inappm', 'users-pr-inappy']"
+                    ]
+                }
+            }
+
     def test_update_quantity_must_be_at_least_2_if_paid_plan(self):
         desired_plan = {"value": PlanName.CODECOV_PRO_YEARLY.value, "quantity": 1}
         response = self._update(
@@ -1109,7 +1214,7 @@ def test_update_sentry_plan_non_sentry_user(
         assert res.json() == {
             "plan": {
                 "value": [
-                    "Invalid value for plan: users-sentrym; must be one of ['users-free', 'users-basic', 'users-pr-inappm', 'users-pr-inappy']"
+                    "Invalid value for plan: users-sentrym; must be one of ['users-basic', 'users-pr-inappm', 'users-pr-inappy']"
                 ]
             }
         }