From 66c7dd3fd239965bba7de67671da6499cfd20167 Mon Sep 17 00:00:00 2001
From: Nora Shapiro <nora.shapiro@sentry.io>
Date: Wed, 28 Aug 2024 12:11:38 -0700
Subject: [PATCH] Update plan service to work with Accounts (#777)

---
 plan/service.py         | 16 +++++++++++-
 plan/tests/test_plan.py | 57 ++++++++++++++++++++++++++++++++++++++++-
 2 files changed, 71 insertions(+), 2 deletions(-)

diff --git a/plan/service.py b/plan/service.py
index 3ab17ca85a..9022901801 100644
--- a/plan/service.py
+++ b/plan/service.py
@@ -41,10 +41,16 @@ def __init__(self, current_org: Owner):
             No value
         """
         self.current_org = current_org
+        self.has_account = False if self.current_org.account is None else True
         if self.current_org.plan not in USER_PLAN_REPRESENTATIONS:
             raise ValueError("Unsupported plan")
         else:
-            self.plan_data = USER_PLAN_REPRESENTATIONS[self.current_org.plan]
+            if self.has_account:
+                self.plan_data = USER_PLAN_REPRESENTATIONS[
+                    self.current_org.account.plan
+                ]
+            else:
+                self.plan_data = USER_PLAN_REPRESENTATIONS[self.current_org.plan]
 
     def update_plan(self, name, user_count: int | None) -> None:
         if name not in USER_PLAN_REPRESENTATIONS:
@@ -75,6 +81,8 @@ def plan_name(self) -> str:
     def plan_user_count(self) -> int:
         if get_config("setup", "enterprise_license"):
             return license_seats()
+        if self.has_account:
+            return self.current_org.account.total_seat_count
         return self.current_org.plan_user_count
 
     @property
@@ -269,6 +277,12 @@ def has_trial_dates(self) -> bool:
     def has_seats_left(self) -> bool:
         if get_config("setup", "enterprise_license"):
             return enterprise_has_seats_left()
+        if self.has_account:
+            # edge case: IF the User is already a plan_activated_user on any of the Orgs in the Account,
+            # AND their Account is at capacity,
+            # AND they try to become a plan_activated_user on another Org in the Account,
+            # has_seats_left will evaluate as False even though the User should be allowed to activate on the Org.
+            return self.current_org.account.can_activate_user()
         return (
             self.plan_activated_users is None
             or len(self.plan_activated_users) < self.plan_user_count
diff --git a/plan/tests/test_plan.py b/plan/tests/test_plan.py
index 9b86644996..3f4109db5e 100644
--- a/plan/tests/test_plan.py
+++ b/plan/tests/test_plan.py
@@ -4,9 +4,10 @@
 from django.test import TestCase
 from freezegun import freeze_time
 from pytest import raises
+from shared.django_apps.codecov_auth.tests.factories import AccountsUsersFactory
 
 from codecov.commands.exceptions import ValidationError
-from codecov_auth.tests.factories import OwnerFactory
+from codecov_auth.tests.factories import AccountFactory, OwnerFactory
 from plan.constants import (
     BASIC_PLAN,
     FREE_PLAN,
@@ -363,6 +364,60 @@ def test_plan_service_update_plan_succeeds(self):
         assert current_org.plan == PlanName.TEAM_MONTHLY.value
         assert current_org.plan_user_count == 8
 
+    def test_has_account(self):
+        current_org = OwnerFactory()
+        plan_service = PlanService(current_org=current_org)
+        self.assertFalse(plan_service.has_account)
+
+        current_org.account = AccountFactory()
+        current_org.save()
+        plan_service = PlanService(current_org=current_org)
+        self.assertTrue(plan_service.has_account)
+
+    def test_plan_data_has_account(self):
+        current_org = OwnerFactory(plan=PlanName.BASIC_PLAN_NAME.value)
+        plan_service = PlanService(current_org=current_org)
+        self.assertEqual(plan_service.plan_name, PlanName.BASIC_PLAN_NAME.value)
+
+        current_org.account = AccountFactory(plan=PlanName.CODECOV_PRO_YEARLY.value)
+        current_org.save()
+        plan_service = PlanService(current_org=current_org)
+        self.assertEqual(plan_service.plan_name, PlanName.CODECOV_PRO_YEARLY.value)
+
+    def test_plan_user_count_has_account(self):
+        org = OwnerFactory(plan=PlanName.BASIC_PLAN_NAME.value, plan_user_count=5)
+        account = AccountFactory(
+            plan=PlanName.BASIC_PLAN_NAME.value, plan_seat_count=50, free_seat_count=3
+        )
+
+        plan_service = PlanService(current_org=org)
+        self.assertEqual(plan_service.plan_user_count, 5)
+
+        org.account = account
+        org.save()
+        plan_service = PlanService(current_org=org)
+        self.assertEqual(plan_service.plan_user_count, 53)
+
+    def test_has_seats_left_has_account(self):
+        org = OwnerFactory(
+            plan=PlanName.BASIC_PLAN_NAME.value,
+            plan_user_count=5,
+            plan_activated_users=[1, 2, 3],
+        )
+        account = AccountFactory(
+            plan=PlanName.BASIC_PLAN_NAME.value, plan_seat_count=5, free_seat_count=3
+        )
+        for i in range(8):
+            AccountsUsersFactory(account=account)
+
+        plan_service = PlanService(current_org=org)
+        self.assertEqual(plan_service.has_seats_left, True)
+
+        org.account = account
+        org.save()
+        plan_service = PlanService(current_org=org)
+        self.assertEqual(plan_service.has_seats_left, False)
+
 
 class AvailablePlansBeforeTrial(TestCase):
     """