diff --git a/documentcloud/organizations/models.py b/documentcloud/organizations/models.py index 0af6ac62..a4261b42 100644 --- a/documentcloud/organizations/models.py +++ b/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): diff --git a/documentcloud/organizations/tests/test_models.py b/documentcloud/organizations/tests/test_models.py index 7a5ad0d9..9916230c 100644 --- a/documentcloud/organizations/tests/test_models.py +++ b/documentcloud/organizations/tests/test_models.py @@ -17,7 +17,7 @@ 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, @@ -25,6 +25,7 @@ def ent_json(entitlement, date_update): "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,7 +189,10 @@ 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() @@ -190,18 +200,31 @@ def test_equal_feature_level_tie_breaks_to_first(self): 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() diff --git a/requirements/base.txt b/requirements/base.txt index bec5853c..9c09ccc4 100644 --- a/requirements/base.txt +++ b/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 diff --git a/requirements/local.txt b/requirements/local.txt index e68d3b86..3482d582 100644 --- a/requirements/local.txt +++ b/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 diff --git a/requirements/production.txt b/requirements/production.txt index 030a497d..e9dd130b 100644 --- a/requirements/production.txt +++ b/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