quotas: refactor/expand tests

src/sentry/quotas/base.py

@@ -39,6 +39,33 @@ def __init__(self, **kwargs):
         super(RateLimited, self).__init__(True, **kwargs)
 
 
+class BasicQuota(object):
+    __slots__ = ['key', 'limit', 'window', 'reason_code', 'enforce']
+
+    def __init__(self, key, limit=0, window=60, reason_code=None,
+                 enforce=True):
+        # the key is effectively the unique identifier for enforcing this quota
+        self.key = key
+        # maximum number of events in the given window, 0 indicates "no limit"
+        self.limit = limit
+        # time in seconds that this quota reflects
+        self.window = window
+        # a machine readable string
+        self.reason_code = reason_code
+        # should this quota be hard-enforced (or just tracked)
+        self.enforce = enforce
+
+    def __eq__(self, other):
+        return isinstance(other, BasicQuota) and hash(self) == hash(other)
+
+    def __hash__(self):
+        return hash((self.key, self.limit, self.window, self.reason_code, self.enforce))
+
+    def __repr__(self):
+        return u'<{} key={} limit={} window={}>'.format(
+            type(self).__name__, self.key, self.limit, self.window)
+
+
 class Quota(Service):
     """
     Quotas handle tracking a project's event usage (at a per minute tick) and
@@ -47,12 +74,59 @@ class Quota(Service):
     """
     __all__ = (
         'get_maximum_quota', 'get_organization_quota', 'get_project_quota',
-        'is_rate_limited', 'translate_quota', 'validate',
+        'get_quotas', 'is_rate_limited', 'translate_quota', 'validate',
     )
 
     def __init__(self, **options):
         pass
 
+    def get_actionable_quotas(self, project, key=None):
+        """
+        Return all implemented quotas which are enabled and actionable.
+
+        This simply suppresses any configured quotas which aren't enabled.
+        """
+        return [
+            quota
+            for quota in self.get_quotas(project, key=key)
+            # a zero limit means "no limit", not "reject all"
+            if quota.limit > 0
+            and quota.window > 0
+        ]
+
+    def get_quotas(self, project, key=None):
+        """
+        Return a list of all configured quotas, even ones which aren't
+        enabled.
+        """
+        if key:
+            key.project = project
+        pquota = self.get_project_quota(project)
+        oquota = self.get_organization_quota(project.organization)
+        results = [
+            BasicQuota(
+                key='p:{}'.format(project.id),
+                limit=pquota[0],
+                window=pquota[1],
+                reason_code='project_quota',
+            ),
+            BasicQuota(
+                key='o:{}'.format(project.organization.id),
+                limit=oquota[0],
+                window=oquota[1],
+                reason_code='org_quota',
+            ),
+        ]
+        if key:
+            kquota = self.get_key_quota(key)
+            results.append(BasicQuota(
+                key='k:{}'.format(key.id),
+                limit=kquota[0],
+                window=kquota[1],
+                reason_code='key_quota',
+            ))
+        return results
+
     def is_rate_limited(self, project, key=None):
         return NotRateLimited()
 
@@ -79,7 +153,8 @@ def get_project_quota(self, project):
 
         org = getattr(project, '_organization_cache', None)
         if not org:
-            org = Organization.objects.get_from_cache(id=project.organization_id)
+            org = Organization.objects.get_from_cache(
+                id=project.organization_id)
             project._organization_cache = org
 
         max_quota_share = int(OrganizationOption.objects.get_value(

src/sentry/quotas/redis.py

@@ -13,26 +13,13 @@
 from time import time
 
 from sentry.exceptions import InvalidConfiguration
-from sentry.quotas.base import NotRateLimited, Quota, RateLimited
+from sentry.quotas.base import BasicQuota, NotRateLimited, Quota, RateLimited
 from sentry.utils.redis import get_cluster_from_options, load_script
 
 is_rate_limited = load_script('quotas/is_rate_limited.lua')
 
 
-class BasicRedisQuota(object):
-    __slots__ = ['key', 'limit', 'window', 'reason_code', 'enforce']
-
-    def __init__(self, key, limit=0, window=60, reason_code=None,
-                 enforce=True):
-        self.key = key
-        # maximum number of events in the given window, 0 indicates "no limit"
-        self.limit = limit
-        # time in seconds that this quota reflects
-        self.window = window
-        # a machine readable string
-        self.reason_code = reason_code
-        # should this quota be hard-enforced (or just tracked)
-        self.enforce = enforce
+BasicRedisQuota = BasicQuota
 
 
 class RedisQuota(Quota):
@@ -60,35 +47,6 @@ def __get_redis_key(self, key, timestamp, interval, shift):
             int((timestamp - shift) // interval),
         )
 
-    def get_quotas(self, project, key=None):
-        if key:
-            key.project = project
-        pquota = self.get_project_quota(project)
-        oquota = self.get_organization_quota(project.organization)
-        results = [
-            BasicRedisQuota(
-                key='p:{}'.format(project.id),
-                limit=pquota[0],
-                window=pquota[1],
-                reason_code='project_quota',
-            ),
-            BasicRedisQuota(
-                key='o:{}'.format(project.organization.id),
-                limit=oquota[0],
-                window=oquota[1],
-                reason_code='org_quota',
-            ),
-        ]
-        if key:
-            kquota = self.get_key_quota(key)
-            results.append(BasicRedisQuota(
-                key='k:{}'.format(key.id),
-                limit=kquota[0],
-                window=kquota[1],
-                reason_code='key_quota',
-            ))
-        return results
-
     def get_usage(self, organization_id, quotas, timestamp=None):
         if timestamp is None:
             timestamp = time()
@@ -132,12 +90,7 @@ def is_rate_limited(self, project, key=None, timestamp=None):
         if timestamp is None:
             timestamp = time()
 
-        quotas = [
-            quota
-            for quota in self.get_quotas(project, key=key)
-            # x = (key, limit, interval)
-            if quota.limit > 0  # a zero limit means "no limit", not "reject all"
-        ]
+        quotas = self.get_actionable_quotas(project, key)
 
         # If there are no quotas to actually check, skip the trip to the database.
         if not quotas:

src/sentry/utils/pytest/sentry.py

@@ -112,6 +112,9 @@ def pytest_configure(config):
         'system.url-prefix': 'http://testserver',
     })
 
+    settings.SENTRY_QUOTAS = 'sentry.quotas.redis.RedisQuota'
+    settings.SENTRY_QUOTA_OPTIONS = {'cluster': 'default'}
+
     # django mail uses socket.getfqdn which doesn't play nice if our
     # networking isn't stable
     patcher = mock.patch('socket.getfqdn', return_value='localhost')

tests/integration/tests.py

@@ -20,7 +20,7 @@
 from six import StringIO
 
 from sentry.models import (
-    Group, GroupTagKey, GroupTagValue, Event, TagKey, TagValue
+    Group, GroupTagKey, GroupTagValue, Event, OrganizationOption, TagKey, TagValue
 )
 from sentry.testutils import TestCase, TransactionTestCase
 from sentry.testutils.helpers import get_auth_header
@@ -398,6 +398,25 @@ def test_protocol_v6(self):
 
         assert instance.message == 'hello'
 
+    def test_basic_quotas(self):
+        results = []
+        for n in range(10):
+            results.append(self._postWithHeader(
+                data={'message': 'hello'},
+                key=self.projectkey.public_key,
+                secret=self.projectkey.secret_key,
+                protocol='6',
+            ))
+
+        # 1 per hour
+        OrganizationOption.objects.set_value(
+            self.project.organization, 'sentry:account-rate-limit', 1,
+        )
+
+        assert results[0].status_code == 200
+        for r in results[1:]:
+            assert r.status_code == 429
+
 
 class DepdendencyTest(TestCase):
     def raise_import_error(self, package):

tests/sentry/quotas/test_base.py

@@ -3,12 +3,13 @@
 from __future__ import absolute_import
 
 from sentry.models import OrganizationOption, ProjectKey
-from sentry.quotas.base import Quota
+from sentry.quotas.base import BasicQuota, Quota
 from sentry.testutils import TestCase
 
 
 class QuotaTest(TestCase):
     def setUp(self):
+        super(QuotaTest, self).setUp()
         self.backend = Quota()
 
     def test_get_project_quota(self):
@@ -37,3 +38,28 @@ def test_get_key_quota(self):
         key = ProjectKey.objects.create(
             project=self.project, rate_limit_window=None, rate_limit_count=None)
         assert self.backend.get_key_quota(key) == (0, 0)
+
+    def test_account_limit(self):
+        org = self.create_organization()
+        project = self.create_project(organization=org)
+
+        OrganizationOption.objects.set_value(
+            org, 'sentry:account-rate-limit', 80,
+        )
+
+        with self.options({'system.rate-limit': 0}):
+            quotas = self.backend.get_quotas(project)
+            assert BasicQuota(
+                key='o:{}'.format(org.id),
+                limit=80,
+                window=3600,
+                reason_code='org_quota',
+            ) in quotas
+
+    def test_ignores_disabled_quotas(self):
+        org = self.create_organization()
+        project = self.create_project(organization=org)
+
+        with self.options({'system.rate-limit': 0}):
+            quotas = self.backend.get_actionable_quotas(project)
+            assert len(quotas) == 0