diff --git a/documentcloud/organizations/models.py b/documentcloud/organizations/models.py index 232a9932..a9b1bed8 100644 --- a/documentcloud/organizations/models.py +++ b/documentcloud/organizations/models.py @@ -122,17 +122,14 @@ def _update_resources(self, data, date_update): @staticmethod def _calc_ent_ai_credits(ent_data): - """Compute the AI credit quota for one entitlement entry. + """The AI credit quota one entitlement entry grants: base × quantity. - Uses ent_data["quantity"] (the per-subscription billing quantity sent - by Squarelet) rather than the stale org-level max_users field. + Squarelet shapes its entitlements so this matches the old + `base + extra users × per user` for every quantity; a plan is quantity 1 + and extra credits are a pack line, whose quantity is how many it holds. """ - 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) - quantity = ent_data.get("quantity", 1) - return base + max(0, quantity - minimum) * per_user + base = ent_data["resources"].get("base_ai_credits", 0) + return base * ent_data.get("quantity", 1) @transaction.atomic def merge(self, uuid): diff --git a/documentcloud/organizations/tests/factories.py b/documentcloud/organizations/tests/factories.py index 2a32f279..c20efabf 100644 --- a/documentcloud/organizations/tests/factories.py +++ b/documentcloud/organizations/tests/factories.py @@ -80,7 +80,7 @@ class ProfessionalEntitlementFactory(EntitlementFactory): resources = { "minimum_users": 1, "base_ai_credits": 2000, - "ai_credits_per_user": 0, + "ai_credits_per_user": 2000, "feature_level": 1, } @@ -90,8 +90,19 @@ class OrganizationEntitlementFactory(EntitlementFactory): name = "Organization" resources = { - "minimum_users": 5, + "minimum_users": 1, "base_ai_credits": 5000, - "ai_credits_per_user": 500, + "ai_credits_per_user": 5000, "feature_level": 2, } + + +class CreditPackEntitlementFactory(EntitlementFactory): + """Extra AI credits, sold by quantity alongside a plan""" + + name = "DocumentCloud Credit Pack" + resources = { + "minimum_users": 1, + "base_ai_credits": 500, + "ai_credits_per_user": 500, + } diff --git a/documentcloud/organizations/tests/test_models.py b/documentcloud/organizations/tests/test_models.py index ae84d5f1..f1e58bd5 100644 --- a/documentcloud/organizations/tests/test_models.py +++ b/documentcloud/organizations/tests/test_models.py @@ -8,6 +8,7 @@ from documentcloud.organizations.exceptions import InsufficientAICreditsError from documentcloud.organizations.models import Organization from documentcloud.organizations.tests.factories import ( + CreditPackEntitlementFactory, EntitlementFactory, OrganizationEntitlementFactory, OrganizationFactory, @@ -103,7 +104,7 @@ def _org_data(self, organization, entitlements): def test_two_paid_entitlements_sums_ai_credits(self): """Two paid entitlements: ai_credits_per_month = sum of both""" ent1 = ProfessionalEntitlementFactory() # base_ai_credits=2000, min=1 - ent2 = OrganizationEntitlementFactory() # base_ai_credits=5000, min=5 + ent2 = OrganizationEntitlementFactory() # base_ai_credits=5000, quantity 1 organization = OrganizationFactory() date_update = date(2024, 3, 1) @@ -112,7 +113,7 @@ def test_two_paid_entitlements_sums_ai_credits(self): organization, [ ent_json(ent1, date_update, quantity=1), - ent_json(ent2, date_update, quantity=5), + ent_json(ent2, date_update, quantity=1), ], ) ) @@ -142,7 +143,7 @@ def test_paid_and_grant_entitlement_sums_ai_credits(self): self._org_data( organization, [ - ent_json(paid, date_update, quantity=5), + ent_json(paid, date_update, quantity=1), ent_json(grant, date_update, quantity=1), ], ) @@ -165,7 +166,7 @@ def test_primary_entitlement_is_highest_feature_level(self): organization, [ ent_json(low, date_update, quantity=1), - ent_json(high, date_update, quantity=5), + ent_json(high, date_update, quantity=1), ], ) ) @@ -200,29 +201,34 @@ def test_equal_feature_level_tie_breaks_to_first(self): assert organization.ai_credits_per_month == 300 @pytest.mark.django_db() - 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 + def test_the_quota_is_base_times_quantity(self): + """Squarelet's entitlements are shaped so this is the whole formula.""" + ent = EntitlementFactory(name="Flat", resources={"base_ai_credits": 500}) organization = OrganizationFactory() organization.update_data( - self._org_data(organization, [ent_json(ent, date(2024, 3, 1), quantity=2)]) + self._org_data(organization, [ent_json(ent, date(2024, 3, 1), quantity=3)]) ) organization.refresh_from_db() - # max(0, 2-5) = 0, so just base=5000 - assert organization.ai_credits_per_month == 5000 + assert organization.ai_credits_per_month == 1500 @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 + def test_a_pack_s_quantity_multiplies_its_credits(self): + """A plan is quantity 1; extra credits are a pack line's quantity.""" + ent = OrganizationEntitlementFactory() + pack = CreditPackEntitlementFactory() organization = OrganizationFactory() organization.update_data( - self._org_data(organization, [ent_json(ent, date(2024, 3, 1), quantity=8)]) + self._org_data( + organization, + [ + ent_json(ent, date(2024, 3, 1), quantity=1), + ent_json(pack, date(2024, 3, 1), quantity=3), + ], + ) ) organization.refresh_from_db() - # 5000 + max(0, 8-5)*500 = 5000 + 1500 = 6500 assert organization.ai_credits_per_month == 6500 @pytest.mark.django_db() @@ -242,7 +248,7 @@ def test_multi_entitlement_monthly_restore(self): organization, [ ent_json(ent1, date(2024, 3, 1), quantity=1), - ent_json(ent2, date(2024, 3, 1), quantity=5), + ent_json(ent2, date(2024, 3, 1), quantity=1), ], ) )