diff --git a/server/lib/python/cartodb_geocoder/cartodb_geocoder/quota_service.py b/server/lib/python/cartodb_geocoder/cartodb_geocoder/quota_service.py index 490c0b9..5d61fad 100644 --- a/server/lib/python/cartodb_geocoder/cartodb_geocoder/quota_service.py +++ b/server/lib/python/cartodb_geocoder/cartodb_geocoder/quota_service.py @@ -15,7 +15,8 @@ class QuotaService: user_quota = self.user_service.user_quota() today = date.today() current_used = self.user_service.used_quota_month(today.year, today.month) - return True if (current_used + 1) < user_quota else False + soft_geocoder_limit = self.user_service.soft_geocoder_limit() + return True if soft_geocoder_limit or (current_used + 1) < user_quota else False def increment_geocoder_use(self, amount=1): today = date.today() diff --git a/server/lib/python/cartodb_geocoder/cartodb_geocoder/user_service.py b/server/lib/python/cartodb_geocoder/cartodb_geocoder/user_service.py index b804f14..3e22d48 100644 --- a/server/lib/python/cartodb_geocoder/cartodb_geocoder/user_service.py +++ b/server/lib/python/cartodb_geocoder/cartodb_geocoder/user_service.py @@ -5,6 +5,8 @@ class UserService: """ Class to manage all the user info """ GEOCODING_QUOTA_KEY = "geocoding_quota" + GEOCODING_SOFT_LIMIT_KEY = "soft_geocoder_limit" + REDIS_CONNECTION_KEY = "redis_connection" REDIS_CONNECTION_HOST = "redis_host" REDIS_CONNECTION_PORT = "redis_port" @@ -29,6 +31,11 @@ class UserService: user_quota = self._redis_connection.hget(self.__get_user_redis_key(), self.GEOCODING_QUOTA_KEY) return int(user_quota) if user_quota and int(user_quota) >= 0 else 0 + def soft_geocoder_limit(self): + """ Check what kind of limit the user has """ + soft_limit = self._redis_connection.hget(self.__get_user_redis_key(), self.GEOCODING_SOFT_LIMIT_KEY) + return True if soft_limit == '1' else False + def used_quota_month(self, year, month): """ Recover the used quota for the user in the current month """ # Check for exceptions or redis timeout diff --git a/server/lib/python/cartodb_geocoder/test/test_quota_service.py b/server/lib/python/cartodb_geocoder/test/test_quota_service.py index d9c5a50..d4c673a 100644 --- a/server/lib/python/cartodb_geocoder/test/test_quota_service.py +++ b/server/lib/python/cartodb_geocoder/test/test_quota_service.py @@ -8,14 +8,14 @@ class TestQuotaService(TestCase): def setUp(self): self.fake_redis_connection = MockRedis() + self.fake_redis_connection.hset('geocoder:user_id','geocoding_quota', 100) + self.fake_redis_connection.hset('geocoder:user_id','soft_geocoder_limit', 0) self.qs = quota_service.QuotaService('user_id', 'tx_id', redis_connection = self.fake_redis_connection) def test_should_return_true_if_quota_with_no_use(self): - self.fake_redis_connection.hset('geocoder:user_id','geocoding_quota', 100) assert self.qs.check_user_quota() == True def test_should_return_true_if_quota_is_not_completely_used(self): - self.fake_redis_connection.hset('geocoder:user_id','geocoding_quota', 100) self.fake_redis_connection.hset('geocoder:user_id:201511','tx_id', 10) self.fake_redis_connection.hset('geocoder:user_id:201511','tx_id_2', 10) assert self.qs.check_user_quota() == True @@ -24,4 +24,11 @@ class TestQuotaService(TestCase): self.fake_redis_connection.hset('geocoder:user_id','geocoding_quota', 1) self.fake_redis_connection.hset('geocoder:user_id:201511','tx_id', 10) self.fake_redis_connection.hset('geocoder:user_id:201511','tx_id_2', 10) - assert self.qs.check_user_quota() == False \ No newline at end of file + assert self.qs.check_user_quota() == False + + def test_should_return_true_if_quota_is_surpassed_but_soft_limit_is_enabled(self): + self.fake_redis_connection.hset('geocoder:user_id','geocoding_quota', 1) + self.fake_redis_connection.hset('geocoder:user_id','soft_geocoder_limit', 1) + self.fake_redis_connection.hset('geocoder:user_id:201511','tx_id', 10) + self.fake_redis_connection.hset('geocoder:user_id:201511','tx_id_2', 10) + assert self.qs.check_user_quota() == True \ No newline at end of file