diff --git a/.isort.cfg b/.isort.cfg index ee97e390..50935845 100755 --- a/.isort.cfg +++ b/.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 diff --git a/documentcloud/core/tests.py b/documentcloud/core/tests.py index d38420b9..e6379117 100644 --- a/documentcloud/core/tests.py +++ b/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() diff --git a/documentcloud/documents/models/saved_search.py b/documentcloud/documents/models/saved_search.py index f071a2e2..d0777417 100644 --- a/documentcloud/documents/models/saved_search.py +++ b/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 _ diff --git a/documentcloud/documents/modifications.py b/documentcloud/documents/modifications.py index 1682dc58..cb183fbb 100644 --- a/documentcloud/documents/modifications.py +++ b/documentcloud/documents/modifications.py @@ -1,4 +1,3 @@ -# Standard Library # Django from django.db import transaction diff --git a/documentcloud/organizations/models.py b/documentcloud/organizations/models.py index 131902b6..0af6ac62 100644 --- a/documentcloud/organizations/models.py +++ b/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 diff --git a/documentcloud/organizations/tests/test_models.py b/documentcloud/organizations/tests/test_models.py index d61a4090..7a5ad0d9 100644 --- a/documentcloud/organizations/tests/test_models.py +++ b/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""" diff --git a/requirements/base.txt b/requirements/base.txt index 2660a3cb..bec5853c 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.14 +squarelet-auth==0.1.15 # via -r requirements/base.in stack-data==0.3.0 # via ipython diff --git a/requirements/local.txt b/requirements/local.txt index 0e77b19f..e68d3b86 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.14 +squarelet-auth==0.1.15 # via -r requirements/base.txt stack-data==0.3.0 # via diff --git a/requirements/production.txt b/requirements/production.txt index 5c5da4a4..030a497d 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.14 +squarelet-auth==0.1.15 # via -r requirements/base.txt stack-data==0.3.0 # via diff --git a/tasks.py b/tasks.py index c9aa4e3b..8f999993 100755 --- a/tasks.py +++ b/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"""