Browse Source

Update org throttling to use stripe account info

David Burke 4 years ago
parent
commit
b163a78e9b

+ 0 - 6
djstripe_ext/tests.py

@@ -3,15 +3,9 @@ from django.shortcuts import reverse
 from django.utils import timezone
 from rest_framework.test import APITestCase
 from model_bakery import baker
-from model_bakery.random_gen import gen_slug, gen_datetime, gen_integer
 from glitchtip import test_utils  # pylint: disable=unused-import
 
 
-baker.generators.add("djstripe.fields.StripeIdField", gen_slug)
-baker.generators.add("djstripe.fields.StripeDateTimeField", gen_datetime)
-baker.generators.add("djstripe.fields.StripeQuantumCurrencyAmountField", gen_integer)
-
-
 class SubscriptionAPITestCase(APITestCase):
     def setUp(self):
         self.user = baker.make("users.user")

+ 5 - 0
glitchtip/settings.py

@@ -12,6 +12,7 @@ https://docs.djangoproject.com/en/dev/ref/settings/
 
 import os
 import sys
+import warnings
 import environ
 import sentry_sdk
 from sentry_sdk.integrations.django import DjangoIntegration
@@ -380,3 +381,7 @@ elif TESTING:
 if TESTING:
     CELERY_TASK_ALWAYS_EAGER = True
     STATICFILES_STORAGE = None
+    # https://github.com/evansd/whitenoise/issues/215
+    warnings.filterwarnings(
+        "ignore", message="No directory at", module="whitenoise.base"
+    )

+ 4 - 1
glitchtip/test_utils/generators.py

@@ -1,5 +1,5 @@
 from model_bakery import baker
-from model_bakery.random_gen import gen_slug
+from model_bakery.random_gen import gen_slug, gen_datetime, gen_integer
 
 
 def currency_code():
@@ -8,3 +8,6 @@ def currency_code():
 
 baker.generators.add("organizations.fields.SlugField", gen_slug)
 baker.generators.add("djstripe.fields.StripeCurrencyCodeField", currency_code)
+baker.generators.add("djstripe.fields.StripeIdField", gen_slug)
+baker.generators.add("djstripe.fields.StripeDateTimeField", gen_datetime)
+baker.generators.add("djstripe.fields.StripeQuantumCurrencyAmountField", gen_integer)

+ 19 - 13
organizations_ext/tasks.py

@@ -1,7 +1,5 @@
 from django.conf import settings
-from django.utils import timezone
-from django.db.models import Count, Q
-from dateutil.relativedelta import relativedelta
+from django.db.models import Count, Q, F
 from celery import shared_task
 from .models import Organization
 
@@ -9,25 +7,33 @@ from .models import Organization
 @shared_task
 def set_organization_throttle():
     """ Determine if organization should be throttled """
-    # Currently throttling only happens if billing is enabled and user has no stripe account.
+    # Currently throttling only happens if billing is enabled and user has free plan.
     if settings.BILLING_ENABLED:
-        # Throttle range is 1 month (not 30 days)
-        month_ago = timezone.now() + relativedelta(months=-1)
         events_max = settings.BILLING_FREE_TIER_EVENTS
-        non_subscriber_organizations = Organization.objects.filter(
-            djstripe_customers__isnull=True,
+        free_tier_organizations = Organization.objects.filter(
+            djstripe_customers__subscriptions__plan__amount=0,
+            djstripe_customers__subscriptions__status="active",
         ).annotate(
-            event_count=Count("projects__issue__event", filter=Q(created__gt=month_ago))
+            event_count=Count(
+                "projects__issue__event",
+                filter=Q(
+                    projects__issue__event__created__gte=F(
+                        "djstripe_customers__subscriptions__current_period_start"
+                    )
+                ),
+            )
         )
 
-        non_subscriber_organizations.filter(
+        free_tier_organizations.filter(
             is_accepting_events=True, event_count__gt=events_max
         ).update(is_accepting_events=False)
-        non_subscriber_organizations.filter(
+        free_tier_organizations.filter(
             is_accepting_events=False, event_count__lte=events_max
         ).update(is_accepting_events=True)
 
-        # is_accepting_events is essentially cache and cache invalidation is hard
+        # paid accounts should always be active at this time
         Organization.objects.filter(
-            is_accepting_events=False, djstripe_customers__isnull=False
+            is_accepting_events=False,
+            djstripe_customers__subscriptions__plan__amount__gt=0,
+            djstripe_customers__subscriptions__status="active",
         ).update(is_accepting_events=True)

+ 17 - 1
organizations_ext/tests/test_throttling.py

@@ -9,8 +9,20 @@ from ..tasks import set_organization_throttle
 class OrganizationThrottlingTestCase(TestCase):
     @override_settings(BILLING_FREE_TIER_EVENTS=10)
     def test_non_subscriber_throttling(self):
+        plan = baker.make("djstripe.Plan", active=True, amount=0)
+
         with freeze_time(timezone.datetime(2000, 1, 1)):
             organization = baker.make("organizations_ext.Organization")
+            customer = baker.make(
+                "djstripe.Customer", subscriber=organization, livemode=False
+            )
+            subscription = baker.make(
+                "djstripe.Subscription",
+                customer=customer,
+                livemode=False,
+                plan=plan,
+                status="active",
+            )
             baker.make(
                 "issues.Event", issue__project__organization=organization, _quantity=3
             )
@@ -27,6 +39,10 @@ class OrganizationThrottlingTestCase(TestCase):
 
         with freeze_time(timezone.datetime(2000, 2, 1)):
             # Month should reset throttle
+            subscription.current_period_start = timezone.make_aware(
+                timezone.datetime(2000, 2, 1)
+            )
+            subscription.save()
             set_organization_throttle()
             organization.refresh_from_db()
             self.assertTrue(organization.is_accepting_events)
@@ -37,4 +53,4 @@ class OrganizationThrottlingTestCase(TestCase):
             )
             set_organization_throttle()
             organization.refresh_from_db()
-            self.assertTrue(organization.is_accepting_events)
+            self.assertFalse(organization.is_accepting_events)