update has_seats_left on plan service to work properly with self-hosted and DEC (#748)

Author: nora-codecov

graphql_api/tests/test_plan.py

@@ -113,7 +113,7 @@ def test_plan_user_count_for_enterprise_org(self, mocked_license):
             is_valid=True,
             message=None,
             url="https://codeov.mysite.com",
-            number_allowed_users=50,
+            number_allowed_users=5,
             number_allowed_repos=10,
             expires=datetime.strptime("2020-05-09 00:00:00", "%Y-%m-%d %H:%M:%S"),
             is_trial=False,
@@ -131,6 +131,21 @@ def test_plan_user_count_for_enterprise_org(self, mocked_license):
             plan_user_count=1,
             plan_activated_users=[],
         )
+        for i in range(4):
+            new_owner = OwnerFactory()
+            enterprise_org.plan_activated_users.append(new_owner.ownerid)
+        enterprise_org.save()
+
+        other_org_in_enterprise = OwnerFactory(
+            service="github",
+            plan=PlanName.CODECOV_PRO_YEARLY.value,
+            plan_user_count=1,
+            plan_activated_users=[],
+        )
+        for i in range(4):
+            new_owner = OwnerFactory()
+            other_org_in_enterprise.plan_activated_users.append(new_owner.ownerid)
+        other_org_in_enterprise.save()
 
         query = """{
                     owner(username: "%s") {
@@ -142,5 +157,35 @@ def test_plan_user_count_for_enterprise_org(self, mocked_license):
                 }
                 """ % (enterprise_org.username)
         data = self.gql_request(query, owner=enterprise_org)
-        assert data["owner"]["plan"]["planUserCount"] == 50
-        assert data["owner"]["plan"]["hasSeatsLeft"] == True
+        assert data["owner"]["plan"]["planUserCount"] == 5
+        assert data["owner"]["plan"]["hasSeatsLeft"] == False
+
+    @patch("services.self_hosted.get_current_license")
+    def test_plan_user_count_for_enterprise_org_invaild_license(self, mocked_license):
+        mock_enterprise_license = LicenseInformation(
+            is_valid=False,
+        )
+        mocked_license.return_value = mock_enterprise_license
+        mock_config_helper(
+            self.mocker, configs={"setup.enterprise_license": mock_enterprise_license}
+        )
+
+        enterprise_org = OwnerFactory(
+            username="random-plan-user",
+            service="github",
+            plan=PlanName.CODECOV_PRO_YEARLY.value,
+            plan_user_count=1,
+            plan_activated_users=[],
+        )
+        query = """{
+                        owner(username: "%s") {
+                            plan {
+                                planUserCount
+                                hasSeatsLeft
+                            }
+                        }
+                    }
+                    """ % (enterprise_org.username)
+        data = self.gql_request(query, owner=enterprise_org)
+        assert data["owner"]["plan"]["planUserCount"] == 0
+        assert data["owner"]["plan"]["hasSeatsLeft"] == False

plan/service.py

@@ -21,7 +21,7 @@
     TrialStatus,
 )
 from services import sentry
-from services.self_hosted import license_seats
+from services.self_hosted import enterprise_has_seats_left, license_seats
 from utils.config import get_config
 
 log = logging.getLogger(__name__)
@@ -267,6 +267,8 @@ def has_trial_dates(self) -> bool:
 
     @property
     def has_seats_left(self) -> bool:
+        if get_config("setup", "enterprise_license"):
+            return enterprise_has_seats_left()
         return (
             self.plan_activated_users is None
             or len(self.plan_activated_users) < self.plan_user_count

services/self_hosted.py

@@ -7,6 +7,7 @@
 from django.db.models import F, Func, Q, QuerySet
 from shared.license import get_current_license
 
+from codecov.db import sync_to_async
 from codecov_auth.models import Owner
 from services import ServiceException
 from utils.config import get_config
@@ -44,16 +45,16 @@ def admin_owners() -> QuerySet:
 
 def is_admin_owner(owner: Optional[Owner]) -> bool:
     """
-    Returns true iff the given owner is an admin.
+    Returns true if the given owner is an admin.
     """
     return owner is not None and admin_owners().filter(pk=owner.pk).exists()
 
 
-def activated_owners() -> QuerySet:
-    """
-    Returns all owners that are activated in ANY org's `plan_activated_users`
-    across the entire instance.
-    """
+async def activated_owners_async() -> QuerySet:
+    return await sync_to_async(activated_owner_query)()
+
+
+def activated_owner_query() -> QuerySet:
     owner_ids = (
         Owner.objects.annotate(
             plan_activated_owner_ids=Func(
@@ -64,13 +65,20 @@ def activated_owners() -> QuerySet:
         .values_list("plan_activated_owner_ids", flat=True)
         .distinct()
     )
-
     return Owner.objects.filter(pk__in=owner_ids)
 
 
+def activated_owners() -> QuerySet:
+    """
+    Returns all owners that are activated in ANY org's `plan_activated_users`
+    across the entire instance.
+    """
+    return activated_owner_query()
+
+
 def is_activated_owner(owner: Owner) -> bool:
     """
-    Returns true iff the given owner is activated in this instance.
+    Returns true if the given owner is activated in this instance.
     """
     return activated_owners().filter(pk=owner.pk).exists()
 
@@ -79,13 +87,27 @@ def license_seats() -> int:
     """
     Max number of seats allowed by the current license.
     """
-    license = get_current_license()
-    return license.number_allowed_users or 0
+    enterprise_license = get_current_license()
+    if not enterprise_license.is_valid:
+        return 0
+    return enterprise_license.number_allowed_users or 0
+
+
+async def enterprise_has_seats_left() -> bool:
+    """
+    The activated_owner_query is heavy, so check the license first, only proceed if they have a valid license.
+    """
+    license_seat_count = license_seats()
+    if license_seat_count == 0:
+        return False
+    owners = await activated_owners_async()
+    count = await sync_to_async(owners.count)()
+    return count < license_seat_count
 
 
 def can_activate_owner(owner: Owner) -> bool:
     """
-    Returns true iff there are available seats left for activation.
+    Returns true if there are available seats left for activation.
     """
     if is_activated_owner(owner):
         # user is already activated in at least 1 org

services/tests/test_self_hosted.py

@@ -87,12 +87,14 @@ def test_is_activated_owner(self, activated_owners):
 
     @patch("services.self_hosted.get_current_license")
     def test_license_seats(self, get_current_license):
-        get_current_license.return_value = LicenseInformation(number_allowed_users=123)
+        get_current_license.return_value = LicenseInformation(
+            number_allowed_users=123, is_valid=True
+        )
         assert license_seats() == 123
 
     @patch("services.self_hosted.get_current_license")
     def test_license_seats_not_specified(self, get_current_license):
-        get_current_license.return_value = LicenseInformation()
+        get_current_license.return_value = LicenseInformation(is_valid=True)
         assert license_seats() == 0
 
     @patch("services.self_hosted.activated_owners")