support multiple entitlements

.isort.cfg

@@ -5,6 +5,7 @@ skip=
 skip_glob=*/migrations/*
 known_future_library=future
 known_django=django,celery,rest_framework
+known_third_party=pysolr
 known_first_party=documentcloud
 indent='    '
 multi_line_output=3

documentcloud/core/tests.py

@@ -1,4 +1,5 @@
 # Django
+from django.conf import settings
 from django.contrib.flatpages.models import FlatPage
 from django.contrib.sites.models import Site
 from django.db import transaction
@@ -58,7 +59,7 @@ def sign(self, data):
         token = "testtoken"
         timestamp = int(time.time())
         signature = hmac.new(
-            key="".encode("utf8"),
+            key=settings.MAILGUN_API_KEY.encode("utf8"),
             msg=f"{timestamp}{token}".encode("utf8"),
             digestmod=hashlib.sha256,
         ).hexdigest()

documentcloud/documents/models/saved_search.py

@@ -1,4 +1,3 @@
-# Standard Library
 # Django
 from django.db import models
 from django.utils.translation import gettext_lazy as _

documentcloud/documents/modifications.py

@@ -1,4 +1,3 @@
-# Standard Library
 # Django
 from django.db import transaction
 

documentcloud/organizations/models.py

@@ -91,8 +91,12 @@ def has_admin(self, user):
         return self.users.filter(pk=user.pk, memberships__admin=True).exists()
 
     def _update_resources(self, data, date_update):
-        # calc reqs/month in case it has changed
-        self.ai_credits_per_month = self.calc_ai_credits_per_month(data["max_users"])
+        # 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"]
+        )
 
         # if date update has changed, then this is a monthly restore of the
         # subscription, and we should restore monthly AI credits.  If not, AI credits
@@ -109,8 +113,13 @@ def _update_resources(self, data, date_update):
             self.monthly_ai_credits = self.ai_credits_per_month
             self.date_update = date_update
 
-    def _choose_entitlement(self, entitlements):
-        return max(entitlements, key=lambda e: e["resources"].get("base_ai_credits", 0))
+    @staticmethod
+    def _calc_ent_ai_credits(ent_data, users):
+        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
 
     @transaction.atomic
     def merge(self, uuid):
@@ -145,15 +154,6 @@ def merge(self, uuid):
 
         self.merged = other
 
-    def calc_ai_credits_per_month(self, users):
-        """Calculate how many AI credits an organization gets per month on this plan
-        for a given number of users"""
-        return (
-            self.entitlement.base_ai_credits
-            + (users - self.entitlement.minimum_users)
-            * self.entitlement.ai_credits_per_user
-        )
-
     @transaction.atomic
     def use_ai_credits(self, amount, user_id, note):
         """Try to deduct AI credits from the organization's balance

documentcloud/organizations/tests/test_models.py

@@ -1,14 +1,33 @@
+# Standard Library
+from datetime import date
+
 # Third Party
 import pytest
 
 # DocumentCloud
 from documentcloud.organizations.exceptions import InsufficientAICreditsError
 from documentcloud.organizations.models import Organization
-from documentcloud.organizations.tests.factories import OrganizationFactory
+from documentcloud.organizations.tests.factories import (
+    EntitlementFactory,
+    OrganizationEntitlementFactory,
+    OrganizationFactory,
+    ProfessionalEntitlementFactory,
+)
 from documentcloud.users.models import User
 from documentcloud.users.tests.factories import UserFactory
 
 
+def ent_json(entitlement, date_update):
+    """Helper function for serializing entitlement data"""
+    return {
+        "name": entitlement.name,
+        "slug": entitlement.slug,
+        "description": entitlement.description,
+        "resources": entitlement.resources,
+        "date_update": date_update,
+    }
+
+
 class TestOrganization:
 
     @pytest.mark.django_db()
@@ -66,6 +85,150 @@ def test_merge_fks(self):
         )
 
 
+class TestSquareletUpdateDataMultiEntitlement:
+    """Test cases for update_data with multiple entitlements"""
+
+    def _org_data(self, organization, entitlements, max_users=5):
+        return {
+            "name": organization.name,
+            "slug": organization.slug,
+            "individual": False,
+            "private": False,
+            "entitlements": entitlements,
+            "max_users": max_users,
+            "card": "",
+        }
+
+    @pytest.mark.django_db()
+    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
+        organization = OrganizationFactory()
+        date_update = date(2024, 3, 1)
+
+        organization.update_data(
+            self._org_data(
+                organization,
+                [ent_json(ent1, date_update), ent_json(ent2, date_update)],
+                max_users=5,
+            )
+        )
+        organization.refresh_from_db()
+        # Professional: 2000 + max(0, 5-1)*0 = 2000
+        # Organization: 5000 + max(0, 5-5)*500 = 5000
+        assert organization.ai_credits_per_month == 7000
+        assert organization.monthly_ai_credits == 7000
+
+    @pytest.mark.django_db()
+    def test_paid_and_grant_entitlement_sums_ai_credits(self):
+        """Paid entitlement + grant entitlement: both contribute to total"""
+        paid = OrganizationEntitlementFactory()
+        grant = EntitlementFactory(
+            name="Grant",
+            resources={
+                "minimum_users": 1,
+                "base_ai_credits": 500,
+                "ai_credits_per_user": 0,
+                "feature_level": 0,
+            },
+        )
+        organization = OrganizationFactory()
+        date_update = date(2024, 3, 1)
+
+        organization.update_data(
+            self._org_data(
+                organization,
+                [ent_json(paid, date_update), ent_json(grant, date_update)],
+                max_users=5,
+            )
+        )
+        organization.refresh_from_db()
+        # Organization: 5000, Grant: 500
+        assert organization.ai_credits_per_month == 5500
+        assert organization.monthly_ai_credits == 5500
+
+    @pytest.mark.django_db()
+    def test_primary_entitlement_is_highest_feature_level(self):
+        """org.entitlement FK points to entitlement with highest feature_level"""
+        low = ProfessionalEntitlementFactory()  # feature_level=1
+        high = OrganizationEntitlementFactory()  # feature_level=2
+        organization = OrganizationFactory()
+        date_update = date(2024, 3, 1)
+
+        organization.update_data(
+            self._org_data(
+                organization,
+                [ent_json(low, date_update), ent_json(high, date_update)],
+            )
+        )
+        organization.refresh_from_db()
+        assert organization.entitlement.slug == high.slug
+
+    @pytest.mark.django_db()
+    def test_equal_feature_level_tie_breaks_to_first(self):
+        """Equal feature_level: first entitlement in list wins the FK"""
+        ent1 = EntitlementFactory(
+            name="GrantA",
+            resources={"base_ai_credits": 100, "feature_level": 1},
+        )
+        ent2 = EntitlementFactory(
+            name="GrantB",
+            resources={"base_ai_credits": 200, "feature_level": 1},
+        )
+        organization = OrganizationFactory()
+        date_update = date(2024, 3, 1)
+
+        organization.update_data(
+            self._org_data(
+                organization,
+                [ent_json(ent1, date_update), ent_json(ent2, date_update)],
+            )
+        )
+        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"""
+        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)
+        )
+        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_multi_entitlement_monthly_restore(self):
+        """Monthly restore resets monthly_ai_credits to sum of all entitlements"""
+        ent1 = ProfessionalEntitlementFactory()
+        ent2 = OrganizationEntitlementFactory()
+        organization = OrganizationFactory(
+            entitlement=ent2,
+            date_update=date(2024, 2, 1),
+            ai_credits_per_month=7000,
+            monthly_ai_credits=1000,
+        )
+
+        organization.update_data(
+            self._org_data(
+                organization,
+                [
+                    ent_json(ent1, date(2024, 3, 1)),
+                    ent_json(ent2, date(2024, 3, 1)),
+                ],
+                max_users=5,
+            )
+        )
+        organization.refresh_from_db()
+        assert organization.ai_credits_per_month == 7000
+        assert organization.monthly_ai_credits == 7000
+
+
 class TestOrganizationCollective:
     """Tests for Organization collective resource sharing"""
 

requirements/base.txt

@@ -451,7 +451,7 @@ sqlparse==0.5.4
     # via
     #   django
     #   django-debug-toolbar
-squarelet-auth==0.1.14
+squarelet-auth==0.1.15
     # 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.14
+squarelet-auth==0.1.15
     # 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.14
+squarelet-auth==0.1.15
     # via -r requirements/base.txt
 stack-data==0.3.0
     # via

tasks.py

@@ -117,12 +117,12 @@ def format(c):
     """Format your code"""
     c.run(
         DJANGO_RUN_USER.format(
-            cmd="black documentcloud --exclude migrations && "
+            cmd="sh -c 'black documentcloud --exclude migrations && "
             "black config/urls.py && "
             "black config/settings && "
             "isort documentcloud && "
             "isort config/urls.py && "
-            "isort config/settings"
+            "isort config/settings'"
         )
     )
 
@@ -233,11 +233,13 @@ def download_tesseract_data(c):
     """Download Tesseract data files. Needed to be able to do OCR locally."""
     c.run("cd config/aws/lambda; ./build.sh")
 
+
 @task
 def initialize_minio(c):
     """Initialize Minio bucket and policies for local development"""
     c.run(DJANGO_RUN.format(cmd="python manage.py initialize_minio"))
 
+
 @task
 def deploy_lambdas(c, staging=False):
     """Deploy lambda functions on AWS"""