use quantity from entitlement instead of max_users

documentcloud/organizations/models.py

@@ -92,10 +92,8 @@ def has_admin(self, user):
 
     def _update_resources(self, data, date_update):
         # sum AI credits across all entitlements
-        users = data["max_users"]
         self.ai_credits_per_month = sum(
-            self._calc_ent_ai_credits(ent_data, users)
-            for ent_data in data["entitlements"]
+            self._calc_ent_ai_credits(ent_data) for ent_data in data["entitlements"]
         )
 
         # if date update has changed, then this is a monthly restore of the
@@ -114,12 +112,18 @@ def _update_resources(self, data, date_update):
             self.date_update = date_update
 
     @staticmethod
-    def _calc_ent_ai_credits(ent_data, users):
+    def _calc_ent_ai_credits(ent_data):
+        """Compute the AI credit quota for one entitlement entry.
+
+        Uses ent_data["quantity"] (the per-subscription billing quantity sent
+        by Squarelet) rather than the stale org-level max_users field.
+        """
         r = ent_data["resources"]
         base = r.get("base_ai_credits", 0)
         per_user = r.get("ai_credits_per_user", 0)
         minimum = r.get("minimum_users", 1)
-        return base + max(0, users - minimum) * per_user
+        quantity = ent_data.get("quantity", 1)
+        return base + max(0, quantity - minimum) * per_user
 
     @transaction.atomic
     def merge(self, uuid):

documentcloud/organizations/tests/test_models.py

@@ -17,14 +17,15 @@
 from documentcloud.users.tests.factories import UserFactory
 
 
-def ent_json(entitlement, date_update):
+def ent_json(entitlement, date_update, quantity=1):
     """Helper function for serializing entitlement data"""
     return {
         "name": entitlement.name,
         "slug": entitlement.slug,
         "description": entitlement.description,
         "resources": entitlement.resources,
         "date_update": date_update,
+        "quantity": quantity,
     }
 
 
@@ -88,14 +89,13 @@ def test_merge_fks(self):
 class TestSquareletUpdateDataMultiEntitlement:
     """Test cases for update_data with multiple entitlements"""
 
-    def _org_data(self, organization, entitlements, max_users=5):
+    def _org_data(self, organization, entitlements):
         return {
             "name": organization.name,
             "slug": organization.slug,
             "individual": False,
             "private": False,
             "entitlements": entitlements,
-            "max_users": max_users,
             "card": "",
         }
 
@@ -110,12 +110,14 @@ def test_two_paid_entitlements_sums_ai_credits(self):
         organization.update_data(
             self._org_data(
                 organization,
-                [ent_json(ent1, date_update), ent_json(ent2, date_update)],
-                max_users=5,
+                [
+                    ent_json(ent1, date_update, quantity=1),
+                    ent_json(ent2, date_update, quantity=5),
+                ],
             )
         )
         organization.refresh_from_db()
-        # Professional: 2000 + max(0, 5-1)*0 = 2000
+        # Professional: 2000 + max(0, 1-1)*0 = 2000
         # Organization: 5000 + max(0, 5-5)*500 = 5000
         assert organization.ai_credits_per_month == 7000
         assert organization.monthly_ai_credits == 7000
@@ -139,12 +141,14 @@ def test_paid_and_grant_entitlement_sums_ai_credits(self):
         organization.update_data(
             self._org_data(
                 organization,
-                [ent_json(paid, date_update), ent_json(grant, date_update)],
-                max_users=5,
+                [
+                    ent_json(paid, date_update, quantity=5),
+                    ent_json(grant, date_update, quantity=1),
+                ],
             )
         )
         organization.refresh_from_db()
-        # Organization: 5000, Grant: 500
+        # Organization: 5000 + max(0, 5-5)*500 = 5000, Grant: 500
         assert organization.ai_credits_per_month == 5500
         assert organization.monthly_ai_credits == 5500
 
@@ -159,7 +163,10 @@ def test_primary_entitlement_is_highest_feature_level(self):
         organization.update_data(
             self._org_data(
                 organization,
-                [ent_json(low, date_update), ent_json(high, date_update)],
+                [
+                    ent_json(low, date_update, quantity=1),
+                    ent_json(high, date_update, quantity=5),
+                ],
             )
         )
         organization.refresh_from_db()
@@ -182,26 +189,42 @@ def test_equal_feature_level_tie_breaks_to_first(self):
         organization.update_data(
             self._org_data(
                 organization,
-                [ent_json(ent1, date_update), ent_json(ent2, date_update)],
+                [
+                    ent_json(ent1, date_update, quantity=1),
+                    ent_json(ent2, date_update, quantity=1),
+                ],
             )
         )
         organization.refresh_from_db()
         assert organization.entitlement.slug == ent1.slug
         assert organization.ai_credits_per_month == 300
 
     @pytest.mark.django_db()
-    def test_users_below_minimum_does_not_reduce_base(self):
-        """users < minimum_users: base AI credits are not reduced"""
+    def test_quantity_below_minimum_does_not_reduce_base(self):
+        """quantity < minimum_users: base AI credits are not reduced"""
         ent = OrganizationEntitlementFactory()  # min=5, base=5000, per_user=500
         organization = OrganizationFactory()
 
         organization.update_data(
-            self._org_data(organization, [ent_json(ent, date(2024, 3, 1))], max_users=2)
+            self._org_data(organization, [ent_json(ent, date(2024, 3, 1), quantity=2)])
         )
         organization.refresh_from_db()
         # max(0, 2-5) = 0, so just base=5000
         assert organization.ai_credits_per_month == 5000
 
+    @pytest.mark.django_db()
+    def test_quantity_above_minimum_adds_per_user_credits(self):
+        """quantity > minimum_users: extra quantity adds per-user AI credits"""
+        ent = OrganizationEntitlementFactory()  # min=5, base=5000, per_user=500
+        organization = OrganizationFactory()
+
+        organization.update_data(
+            self._org_data(organization, [ent_json(ent, date(2024, 3, 1), quantity=8)])
+        )
+        organization.refresh_from_db()
+        # 5000 + max(0, 8-5)*500 = 5000 + 1500 = 6500
+        assert organization.ai_credits_per_month == 6500
+
     @pytest.mark.django_db()
     def test_multi_entitlement_monthly_restore(self):
         """Monthly restore resets monthly_ai_credits to sum of all entitlements"""
@@ -218,10 +241,9 @@ def test_multi_entitlement_monthly_restore(self):
             self._org_data(
                 organization,
                 [
-                    ent_json(ent1, date(2024, 3, 1)),
-                    ent_json(ent2, date(2024, 3, 1)),
+                    ent_json(ent1, date(2024, 3, 1), quantity=1),
+                    ent_json(ent2, date(2024, 3, 1), quantity=5),
                 ],
-                max_users=5,
             )
         )
         organization.refresh_from_db()

requirements/base.txt

@@ -451,7 +451,7 @@ sqlparse==0.5.4
     # via
     #   django
     #   django-debug-toolbar
-squarelet-auth==0.1.15
+squarelet-auth==0.1.16
     # via -r requirements/base.in
 stack-data==0.3.0
     # via ipython

requirements/local.txt

@@ -786,7 +786,7 @@ sqlparse==0.5.4
     #   -r requirements/base.txt
     #   django
     #   django-debug-toolbar
-squarelet-auth==0.1.15
+squarelet-auth==0.1.16
     # via -r requirements/base.txt
 stack-data==0.3.0
     # via

requirements/production.txt

@@ -622,7 +622,7 @@ sqlparse==0.5.4
     #   -r requirements/base.txt
     #   django
     #   django-debug-toolbar
-squarelet-auth==0.1.15
+squarelet-auth==0.1.16
     # via -r requirements/base.txt
 stack-data==0.3.0
     # via