86 lines
3.1 KiB
Python
86 lines
3.1 KiB
Python
from unittest import TestCase
|
|
from mockredis import MockRedis
|
|
from ..test_helper import *
|
|
from cartodb_services.metrics.quota import QuotaChecker
|
|
from cartodb_services.metrics import RoutingConfig
|
|
from datetime import datetime
|
|
|
|
|
|
class RoutingConfigMock(object):
|
|
|
|
def __init__(self, **kwargs):
|
|
self.__dict__ = kwargs
|
|
|
|
|
|
class TestQuotaChecker(TestCase):
|
|
|
|
def setUp(self):
|
|
self.username = 'my_test_user'
|
|
self.period_end_date = datetime.today()
|
|
self.service_type = 'routing_mapzen'
|
|
self.redis_key = 'user:{0}:{1}:success_responses:{2}{3}'.format(
|
|
self.username,
|
|
self.service_type,
|
|
self.period_end_date.year,
|
|
self.period_end_date.strftime('%m')
|
|
)
|
|
|
|
def test_routing_quota_check_passes_when_enough_quota(self):
|
|
user_service_config = RoutingConfigMock(
|
|
username = self.username,
|
|
organization = None,
|
|
service_type = self.service_type,
|
|
monthly_quota = 1000,
|
|
period_end_date = datetime.today(),
|
|
soft_limit = False
|
|
)
|
|
redis_conn = MockRedis()
|
|
redis_conn.zincrby(self.redis_key, self.period_end_date.strftime('%d'), 999)
|
|
assert QuotaChecker(user_service_config, redis_conn).check() == True
|
|
|
|
def test_routing_quota_check_fails_when_quota_exhausted(self):
|
|
user_service_config = RoutingConfigMock(
|
|
username = self.username,
|
|
organization = None,
|
|
service_type = self.service_type,
|
|
monthly_quota = 1000,
|
|
period_end_date = datetime.today(),
|
|
soft_limit = False
|
|
)
|
|
redis_conn = MockRedis()
|
|
redis_conn.zincrby(self.redis_key, self.period_end_date.strftime('%d'), 1001)
|
|
checker = QuotaChecker(user_service_config, redis_conn)
|
|
assert checker.check() == False
|
|
|
|
def test_routing_quota_check_fails_right_in_the_limit(self):
|
|
"""
|
|
I have 1000 credits and I just spent 1000 today. I should not pass
|
|
the check to perform the 1001th routing operation.
|
|
"""
|
|
user_service_config = RoutingConfigMock(
|
|
username = self.username,
|
|
organization = None,
|
|
service_type = self.service_type,
|
|
monthly_quota = 1000,
|
|
period_end_date = datetime.today(),
|
|
soft_limit = False
|
|
)
|
|
redis_conn = MockRedis()
|
|
redis_conn.zincrby(self.redis_key, self.period_end_date.strftime('%d'), 1000)
|
|
checker = QuotaChecker(user_service_config, redis_conn)
|
|
assert checker.check() == False
|
|
|
|
def test_routing_quota_check_passes_if_no_quota_but_soft_limit(self):
|
|
user_service_config = RoutingConfigMock(
|
|
username = self.username,
|
|
organization = None,
|
|
service_type = self.service_type,
|
|
monthly_quota = 1000,
|
|
period_end_date = datetime.today(),
|
|
soft_limit = True
|
|
)
|
|
redis_conn = MockRedis()
|
|
redis_conn.zincrby(self.redis_key, self.period_end_date.strftime('%d'), 1001)
|
|
checker = QuotaChecker(user_service_config, redis_conn)
|
|
assert checker.check() == True
|