Merge pull request #29 from CartoDB/extract_quota_logic

Extracted redis connection and quota check to be reusable

.gitignore

@@ -1,2 +1,2 @@
-*.DS_Store
+.DS_Store
 *.pyc

server/extension/sql/0.0.1/10_config.sql

@@ -1,38 +0,0 @@
---
--- This extension has its own table for configurations.
---
--- The table and the function are considered to be private and therefore
--- no permissions are granted for any other user but the creator/geocoder.
-
-CREATE TABLE IF NOT EXISTS cdb_geocoder_server._config ( KEY TEXT PRIMARY KEY, VALUE JSON NOT NULL );
-
--- Needed to dump config in backups
--- This can only be called from an SQL script executed by CREATE EXTENSION
-SELECT pg_catalog.pg_extension_config_dump('cdb_geocoder_server._config', '');
-
-
-CREATE OR REPLACE FUNCTION cdb_geocoder_server._config_set(key text, value JSON)
-RETURNS VOID AS $$
-BEGIN
-    PERFORM cdb_geocoder_server._config_remove(key);
-    EXECUTE 'INSERT INTO cdb_geocoder_server._config (KEY, VALUE) VALUES ($1, $2);' USING key, value;
-END
-$$ LANGUAGE PLPGSQL VOLATILE;
-
-
-CREATE OR REPLACE FUNCTION cdb_geocoder_server._config_remove(key text)
-RETURNS VOID AS $$
-BEGIN
-    EXECUTE 'DELETE FROM cdb_geocoder_server._config WHERE KEY = $1;' USING key;
-END
-$$ LANGUAGE PLPGSQL VOLATILE;
-
-CREATE OR REPLACE FUNCTION cdb_geocoder_server._config_get(key text)
-    RETURNS JSON AS $$
-DECLARE
-    value JSON;
-BEGIN
-    EXECUTE 'SELECT VALUE FROM cdb_geocoder_server._config WHERE KEY = $1;' INTO value USING key;
-    RETURN value;
-END
-$$ LANGUAGE PLPGSQL STABLE;

server/extension/sql/0.0.1/10_redis_helper.sql

@@ -0,0 +1,44 @@
+CREATE TYPE cdb_geocoder_server._redis_conf_params AS (
+    sentinel_host text,
+    sentinel_port int,
+    sentinel_master_id text,
+    redis_db text,
+    timeout float
+);
+
+-- Get the Redis configuration from the _conf table --
+CREATE OR REPLACE FUNCTION cdb_geocoder_server._get_redis_conf()
+RETURNS cdb_geocoder_server._redis_conf_params AS $$
+    conf = plpy.execute("SELECT cartodb.CDB_Conf_GetConf('redis_conf') conf")[0]['conf']
+    if conf is None:
+      plpy.error("There is no redis configuration defined")
+    else:
+      import json
+      params = json.loads(conf)
+      return {
+        "sentinel_host": params['sentinel_host'],
+        "sentinel_port": params['sentinel_port'],
+        "sentinel_master_id": params['sentinel_master_id'],
+        "timeout": params['timeout'],
+        "redis_db": params['redis_db']
+      }
+$$ LANGUAGE plpythonu;
+
+-- Get the connection to redis from cache or create a new one
+CREATE OR REPLACE FUNCTION cdb_geocoder_server._connect_to_redis(user_id text)
+RETURNS boolean AS $$
+  if user_id in GD and 'redis_connection' in GD[user_id]:
+    return False
+  else:
+    from cartodb_geocoder import redis_helper
+    config_params = plpy.execute("""select c.sentinel_host, c.sentinel_port,
+        c.sentinel_master_id, c.timeout, c.redis_db
+        from cdb_geocoder_server._get_redis_conf() c;""")[0]
+    redis_connection = redis_helper.RedisHelper(config_params['sentinel_host'],
+        config_params['sentinel_port'],
+        config_params['sentinel_master_id'],
+        timeout=config_params['timeout'],
+        redis_db=config_params['redis_db']).redis_connection()
+    GD[user_id] = {'redis_connection': redis_connection}
+    return True
+$$ LANGUAGE plpythonu;

server/extension/sql/0.0.1/30_admin0.sql

@@ -1,6 +1,6 @@
 -- Interface of the server extension
 
-CREATE OR REPLACE FUNCTION cdb_geocoder_server.geocode_admin0_polygon(user_id name, tx_id bigint, country_name text)
+CREATE OR REPLACE FUNCTION cdb_geocoder_server.geocode_admin0_polygon(user_id name, user_config_data JSON, geocoder_config_data JSON, country_name text)
 RETURNS Geometry AS $$
     plpy.debug('Entering geocode_admin0_polygons')
     plpy.debug('user_id = %s' % user_id)

server/lib/python/cartodb_geocoder/cartodb_geocoder/config_helper.py

@@ -0,0 +1,104 @@
+import json
+
+class ConfigException(Exception):
+    pass
+
+class UserConfig:
+
+  USER_CONFIG_KEYS = ['is_organization', 'entity_name']
+
+  def __init__(self, user_config_json, db_user_id = None):
+    config = json.loads(user_config_json)
+    filtered_config = { key: config[key] for key in self.USER_CONFIG_KEYS if key in config.keys() }
+    self.__check_config(filtered_config)
+    self.__parse_config(filtered_config)
+    self._user_id = self.__extract_uuid(db_user_id)
+
+  def __check_config(self, filtered_config):
+    if len(filtered_config.keys()) != len(self.USER_CONFIG_KEYS):
+      raise ConfigException("Passed user configuration is not correct, check it please")
+
+    return True
+
+  @property
+  def is_organization(self):
+    return self._is_organization
+
+  @property
+  def entity_name(self):
+    return self._entity_name
+
+  @property
+  def user_id(self):
+    return self._user_id
+
+  def __parse_config(self, filtered_config):
+    self._is_organization = filtered_config['is_organization']
+    self._entity_name = filtered_config['entity_name']
+
+  def __extract_uuid(self, db_user_id):
+    # Format: development_cartodb_user_<UUID>
+    return db_user_id.split('_')[-1]
+
+class GeocoderConfig:
+
+  GEOCODER_CONFIG_KEYS = ['street_geocoder_provider', 'google_maps_private_key',
+    'nokia_monthly_quota', 'nokia_soft_geocoder_limit']
+  NOKIA_GEOCODER_MANDATORY_KEYS = ['nokia_monthly_quota', 'nokia_soft_geocoder_limit']
+  NOKIA_GEOCODER = 'nokia'
+  NOKIA_QUOTA_KEY = 'nokia_monthly_quota'
+  NOKIA_SOFT_LIMIT_KEY = 'nokia_soft_geocoder_limit'
+  GOOGLE_GEOCODER = 'google'
+  GEOCODER_TYPE = 'street_geocoder_provider'
+  GOOGLE_GEOCODER_API_KEY = 'google_maps_private_key'
+
+  def __init__(self, geocoder_config_json):
+    config = json.loads(geocoder_config_json)
+    filtered_config = { key: config[key] for key in self.GEOCODER_CONFIG_KEYS if key in config.keys() }
+    self.__check_config(filtered_config)
+    self.__parse_config(filtered_config)
+
+  def __check_config(self, filtered_config):
+    if filtered_config[self.GEOCODER_TYPE].lower() == self.NOKIA_GEOCODER:
+      if not set(self.NOKIA_GEOCODER_MANDATORY_KEYS).issubset(set(filtered_config.keys())):
+        raise ConfigException("""Nokia geocoder need the mandatory parameters 'nokia_monthly_quota' and 'nokia_soft_geocoder_limit'""")
+    elif filtered_config[self.GEOCODER_TYPE].lower() == self.GOOGLE_GEOCODER:
+      if self.GOOGLE_GEOCODER_API_KEY not in filtered_config.keys():
+        raise ConfigException("Google geocoder need the mandatory parameter 'google_maps_private_key'")
+
+    return True
+
+  def __parse_config(self, filtered_config):
+    self._geocoder_type = filtered_config[self.GEOCODER_TYPE].lower()
+    self._google_maps_private_key = None
+    self._nokia_monthly_quota = 0
+    self._nokia_soft_geocoder_limit = False
+    if self.GOOGLE_GEOCODER == self._geocoder_type:
+      self._google_maps_private_key = filtered_config[self.GOOGLE_GEOCODER_API_KEY]
+    elif self.NOKIA_GEOCODER == self._geocoder_type:
+      self._nokia_monthly_quota = filtered_config[self.NOKIA_QUOTA_KEY]
+      self._nokia_soft_geocoder_limit = filtered_config[self.NOKIA_SOFT_LIMIT_KEY]
+
+  @property
+  def service_type(self):
+    return self._geocoder_type
+
+  @property
+  def nokia_geocoder(self):
+    return self._geocoder_type == self.NOKIA_GEOCODER
+
+  @property
+  def google_geocoder(self):
+    return self._geocoder_type == self.GOOGLE_GEOCODER
+
+  @property
+  def google_api_key(self):
+    return self._google_maps_private_key
+
+  @property
+  def nokia_monthly_quota(self):
+    return self._nokia_monthly_quota
+
+  @property
+  def nokia_soft_limit(self):
+    return self._nokia_soft_geocoder_limit

server/lib/python/cartodb_geocoder/cartodb_geocoder/quota_service.py

@@ -1,27 +1,31 @@
-import redis
 import user_service
+import config_helper
 from datetime import date
 
 class QuotaService:
     """ Class to manage all the quota operation for the Geocoder SQL API Extension """
 
-    def __init__(self, user_id, transaction_id, **kwargs):
-        self._user_service = user_service.UserService(user_id, **kwargs)
-        self.transaction_id = transaction_id
+    def __init__(self, user_config, geocoder_config, redis_connection):
+        self._user_config = user_config
+        self._geocoder_config = geocoder_config
+        self._user_service = user_service.UserService(self._user_config,
+            self._geocoder_config.service_type, redis_connection)
 
     def check_user_quota(self):
         """ Check if the current user quota surpasses the current quota """
-        # TODO We need to add the hard/soft limit flag for the geocoder
-        user_quota = self.user_service.user_quota()
+        # We don't have quota check for google geocoder
+        if self._geocoder_config.google_geocoder:
+            return True
+
+        user_quota = self._geocoder_config.nokia_monthly_quota
         today = date.today()
-        current_used = self.user_service.used_quota_month(today.year, today.month)
-        soft_geocoder_limit = self.user_service.soft_geocoder_limit()
-        return True if soft_geocoder_limit or (current_used + 1) < user_quota else False
+        service_type = self._geocoder_config.service_type
+        current_used = self._user_service.used_quota(service_type, today.year, today.month)
+        soft_geocoder_limit = self._geocoder_config.nokia_soft_limit
+
+        print "User quota: {0} --- current_used: {1} --- limit: {2}".format(user_quota, current_used, soft_geocoder_limit)
+
+        return True if soft_geocoder_limit or current_used <= user_quota else False
 
     def increment_geocoder_use(self, amount=1):
-        today = date.today()
-        self.user_service.increment_geocoder_use(today.year, today.month, self.transaction_id)
-
-    @property
-    def user_service(self):
-        return self._user_service
+        self._user_service.increment_service_use(self._geocoder_config.service_type)

server/lib/python/cartodb_geocoder/cartodb_geocoder/redis_helper.py

@@ -0,0 +1,20 @@
+from redis.sentinel import Sentinel
+
+class RedisHelper:
+
+  REDIS_DEFAULT_USER_DB = 5
+  REDIS_TIMEOUT = 2 #seconds
+
+  def __init__(self, sentinel_host, sentinel_port, sentinel_master_id, redis_db=REDIS_DEFAULT_USER_DB, **kwargs):
+    self.sentinel_host = sentinel_host
+    self.sentinel_port = sentinel_port
+    self.sentinel_master_id = sentinel_master_id
+    self.timeout = kwargs['timeout'] if 'timeout' in kwargs else REDIS_DEFAULT_TIMEOUT
+    self.redis_db = redis_db
+
+  def redis_connection(self):
+    return self.__create_redis_connection()
+
+  def __create_redis_connection(self):
+    sentinel = Sentinel([(self.sentinel_host, 26379)], socket_timeout=self.timeout)
+    return sentinel.master_for(self.sentinel_master_id, socket_timeout=self.timeout, db=self.redis_db)

server/lib/python/cartodb_geocoder/cartodb_geocoder/user_service.py

@@ -1,4 +1,4 @@
-import redis
+import redis_helper
 from datetime import date
 
 class UserService:
@@ -12,68 +12,53 @@ class UserService:
   REDIS_CONNECTION_PORT = "redis_port"
   REDIS_CONNECTION_DB = "redis_db"
 
-  REDIS_DEFAULT_USER_DB = 5
-  REDIS_DEFAULT_HOST = 'localhost'
-  REDIS_DEFAULT_PORT = 6379
+  def __init__(self, user_config, service_type, redis_connection):
+    self.user_config = user_config
+    self.service_type = service_type
+    self._redis_connection = redis_connection
 
-  def __init__(self, user_id, **kwargs):
-    self.user_id = user_id
-    if self.REDIS_CONNECTION_KEY in kwargs:
-      self._redis_connection = self.__get_redis_connection(redis_connection=kwargs[self.REDIS_CONNECTION_KEY])
-    else:
-      if self.REDIS_CONNECTION_HOST not in kwargs:
-        raise Exception("You have to provide redis configuration")
-      redis_config = self.__build_redis_config(kwargs)
-      self._redis_connection = self.__get_redis_connection(redis_config = redis_config)
-
-  def user_quota(self):
-      # Check for exceptions or redis timeout
-      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):
+  def used_quota(self, service_type, year, month, day=None):
       """ Recover the used quota for the user in the current month """
-      # Check for exceptions or redis timeout
-      current_used = 0
-      for _, value in self._redis_connection.hscan_iter(self.__get_month_redis_key(year,month)):
-          current_used += int(value)
-      return current_used
+      redis_key_data = self.__get_redis_key(service_type, year, month, day)
+      current_use = self._redis_connection.hget(redis_key_data['redis_name'], redis_key_data['redis_key'])
+      return int(current_use) if current_use else 0
 
-  def increment_geocoder_use(self, year, month, key, amount=1):
-      # TODO Manage exceptions or timeout
-      self._redis_connection.hincrby(self.__get_month_redis_key(year, month),key,amount)
+  def increment_service_use(self, service_type, date=date.today(), amount=1):
+      """ Increment the services uses in monthly and daily basis"""
+      self.__increment_monthly_uses(date, service_type, amount)
+      self.__increment_daily_uses(date, service_type, amount)
 
-  @property
-  def redis_connection(self):
-      return self._redis_connection
+  # Private functions
 
-  def __get_redis_connection(self, redis_connection=None, redis_config=None):
-      if redis_connection:
-          conn = redis_connection
-      else:
-          conn = self.__create_redis_connection(redis_config)
+  def __increment_monthly_uses(self, date, service_type, amount):
+    redis_key_data = self.__get_redis_key(service_type, date.year, date.month)
+    self._redis_connection.hincrby(redis_key_data['redis_name'],redis_key_data['redis_key'],amount)
 
-      return conn
+  def __increment_daily_uses(self, date, service_type, amount):
+    redis_key_data = self.__get_redis_key(service_type, date.year, date.month, date.day)
+    self._redis_connection.hincrby(redis_key_data['redis_name'],redis_key_data['redis_key'],amount)
 
-  def __create_redis_connection(self, redis_config):
-      pool = redis.ConnectionPool(host=redis_config['host'], port=redis_config['port'], db=redis_config['db'])
-      conn = redis.Redis(connection_pool=pool)
-      return conn
+  def __get_redis_key(self, service_type, year, month, day=None):
+    redis_name = self.__parse_redis_name(service_type,day)
+    redis_key = self.__parse_redis_key(year,month,day)
 
-  def __build_redis_config(self, config):
-      redis_host = config[self.REDIS_CONNECTION_HOST] if self.REDIS_CONNECTION_HOST in config else self.REDIS_DEFAULT_HOST
-      redis_port = config[self.REDIS_CONNECTION_PORT] if self.REDIS_CONNECTION_PORT in config else self.REDIS_DEFAULT_PORT
-      redis_db = config[self.REDIS_CONNECTION_DB] if self.REDIS_CONNECTION_DB in config else self.REDIS_DEFAULT_USER_DB
-      return {'host': redis_host, 'port': redis_port, 'db': redis_db}
+    return {'redis_name': redis_name, 'redis_key': redis_key}
 
-  def __get_month_redis_key(self, year, month):
-      today = date.today()
-      return "geocoder:{0}:{1}{2}".format(self.user_id, year, month)
+  def __parse_redis_name(self,service_type, day=None):
+    prefix = "org" if self.user_config.is_organization else "user"
+    dated_key = "used_quota_day" if day else "used_quota_month"
+    redis_name = "{0}:{1}:{2}:{3}".format(
+      prefix, self.user_config.entity_name, service_type, dated_key
+    )
+    if self.user_config.is_organization and day:
+      redis_name = "{0}:{1}".format(redis_name, self.user_config.user_id)
 
-  def __get_user_redis_key(self):
-      return "geocoder:{0}".format(self.user_id)
+    return redis_name
+
+  def __parse_redis_key(self,year,month,day=None):
+    if day:
+      redis_key = "{0}_{1}_{2}".format(year,month,day)
+    else:
+      redis_key = "{0}_{1}".format(year,month)
+
+    return redis_key

server/lib/python/cartodb_geocoder/test/test_config_helper.py

@@ -0,0 +1,48 @@
+from cartodb_geocoder import config_helper
+from unittest import TestCase
+from nose.tools import assert_raises
+
+
+class TestConfigHelper(TestCase):
+
+  def test_should_return_list_of_user_config_if_its_ok(self):
+    user_config_json = '{"is_organization": false, "entity_name": "test_user"}'
+    user_config = config_helper.UserConfig(user_config_json, 'development_cartodb_user_UUID')
+    assert user_config.is_organization == False
+    assert user_config.entity_name == 'test_user'
+    assert user_config.user_id == 'UUID'
+
+  def test_should_return_raise_config_exception_if_not_ok(self):
+    user_config_json = '{"is_organization": "false"}'
+    assert_raises(config_helper.ConfigException, config_helper.UserConfig, user_config_json)
+
+  def test_should_return_raise_config_exception_if_empty(self):
+    user_config_json = '{}'
+    assert_raises(config_helper.ConfigException, config_helper.UserConfig, user_config_json)
+
+  def test_should_return_list_of_nokia_geocoder_config_if_its_ok(self):
+    geocoder_config_json = """{"street_geocoder_provider": "Nokia",
+      "nokia_monthly_quota": 100, "nokia_soft_geocoder_limit": false}"""
+    geocoder_config = config_helper.GeocoderConfig(geocoder_config_json)
+    assert geocoder_config.nokia_geocoder == True
+    assert geocoder_config.nokia_monthly_quota == 100
+    assert geocoder_config.nokia_soft_limit == False
+
+  def test_should_raise_configuration_exception_when_missing_nokia_geocoder_parameters(self):
+    geocoder_config_json = '{"street_geocoder_provider": "NokiA", "nokia_monthly_quota": "100"}'
+    assert_raises(config_helper.ConfigException, config_helper.GeocoderConfig, geocoder_config_json)
+
+  def test_should_raise_configuration_exception_when_missing_nokia_geocoder_parameters_2(self):
+    geocoder_config_json = '{"street_geocoder_provider": "NoKia", "nokia_soft_geocoder_limit": "false"}'
+    assert_raises(config_helper.ConfigException, config_helper.GeocoderConfig, geocoder_config_json)
+
+  def test_should_return_list_of_google_geocoder_config_if_its_ok(self):
+    geocoder_config_json = """{"street_geocoder_provider": "gOOgle",
+      "google_maps_private_key": "sdasdasda"}"""
+    geocoder_config = config_helper.GeocoderConfig(geocoder_config_json)
+    assert geocoder_config.google_geocoder == True
+    assert geocoder_config.google_api_key ==  'sdasdasda'
+
+  def test_should_raise_configuration_exception_when_missing_google_api_key(self):
+    geocoder_config_json = '{"street_geocoder_provider": "google"}'
+    assert_raises(config_helper.ConfigException, config_helper.GeocoderConfig, geocoder_config_json)

server/lib/python/cartodb_geocoder/test/test_quota_service.py

@@ -1,34 +1,85 @@
 from mockredis import MockRedis
 from cartodb_geocoder import quota_service
+from cartodb_geocoder import config_helper
 from unittest import TestCase
 from nose.tools import assert_raises
+from datetime import datetime
 
 
 class TestQuotaService(TestCase):
 
+  # single user
+  #   user:<username>:<service>:used_quota_month:year_month
+  #   user:<username>:<service>:used_quota_day:year_month_day
+  # organization user
+  #   org:<orgname>:<service>:used_quota_month:year_month
+  #   org:<orgname>:<service>:<uuid>:used_quota_day:year_month_day
+
   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):
-    assert self.qs.check_user_quota() == True
+  def test_should_return_true_if_user_quota_with_no_use(self):
+    qs = self.__build_quota_service()
+    assert 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:201511','tx_id', 10)
-    self.fake_redis_connection.hset('geocoder:user_id:201511','tx_id_2', 10)
-    assert self.qs.check_user_quota() == True
+  def test_should_return_true_if_org_quota_with_no_use(self):
+    qs = self.__build_quota_service(organization=True)
+    assert qs.check_user_quota() == True
 
-  def test_should_return_false_if_quota_is_surpassed(self):
-    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
+  def test_should_return_true_if_user_quota_is_not_completely_used(self):
+    qs = self.__build_quota_service()
+    self.__increment_geocoder_uses('test_user', '20151111')
+    assert qs.check_user_quota() == True
 
-  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
+  def test_should_return_true_if_org_quota_is_not_completely_used(self):
+    qs = self.__build_quota_service(organization=True)
+    self.__increment_geocoder_uses('test_user', '20151111', org=True)
+    assert qs.check_user_quota() == True
+
+  def test_should_return_false_if_user_quota_is_surpassed(self):
+    qs = self.__build_quota_service(quota = 1, soft_limit=False)
+    self.__increment_geocoder_uses('test_user', '20151111')
+    assert qs.check_user_quota() == False
+
+  def test_should_return_false_if_org_quota_is_surpassed(self):
+    qs = self.__build_quota_service(organization=True, quota=1)
+    self.__increment_geocoder_uses('test_user', '20151111', org=True)
+    assert qs.check_user_quota() == False
+
+  def test_should_return_true_if_user_quota_is_surpassed_but_soft_limit_is_enabled(self):
+    qs = self.__build_quota_service(quota=1, soft_limit=True)
+    self.__increment_geocoder_uses('test_user', '20151111')
+    assert qs.check_user_quota() == True
+
+  def test_should_return_true_if_org_quota_is_surpassed_but_soft_limit_is_enabled(self):
+    qs = self.__build_quota_service(organization=True, quota=1, soft_limit=True)
+    self.__increment_geocoder_uses('test_user', '20151111', org=True)
+    assert qs.check_user_quota() == True
+
+  def test_should_check_user_increment_and_quota_check_correctly(self):
+    qs = self.__build_quota_service(quota=2, soft_limit=False)
+    qs.increment_geocoder_use()
+    assert qs.check_user_quota() == True
+
+  def test_should_check_org_increment_and_quota_check_correctly(self):
+    qs = self.__build_quota_service(organization=True, quota=2, soft_limit=False)
+    qs.increment_geocoder_use()
+    assert qs.check_user_quota() == True
+
+  def __build_quota_service(self, quota=100, service='nokia', organization=False, soft_limit=False):
+    is_organization = 'true' if organization else 'false'
+    has_soft_limit = 'true' if soft_limit else 'false'
+    user_config_json = '{{"is_organization": {0}, "entity_name": "test_user"}}'.format(is_organization)
+    geocoder_config_json = """{{"street_geocoder_provider": "{0}","nokia_monthly_quota": {1},
+      "nokia_soft_geocoder_limit": {2}}}""".format(service, quota, has_soft_limit)
+    user_config = config_helper.UserConfig(user_config_json, 'user_id')
+    geocoder_config = config_helper.GeocoderConfig(geocoder_config_json)
+
+    return quota_service.QuotaService(user_config, geocoder_config, redis_connection = self.fake_redis_connection)
+
+  def __increment_geocoder_uses(self, entity_name, date_string, service='nokia', amount=20, org=False):
+    prefix = 'org' if org else 'user'
+    date = datetime.strptime(date_string, "%Y%m%d")
+    redis_name = "{0}:{1}:{2}:used_quota_month".format(prefix, entity_name, service)
+    redis_key_month = "{0}_{1}".format(date.year, date.month)
+    self.fake_redis_connection.hset(redis_name, redis_key_month, amount)

server/lib/python/cartodb_geocoder/test/test_user_service.py

@@ -1,37 +1,63 @@
 from mockredis import MockRedis
 from cartodb_geocoder import user_service
+from cartodb_geocoder import config_helper
+from datetime import datetime
 from unittest import TestCase
 from nose.tools import assert_raises
 
 
 class TestUserService(TestCase):
 
+    NOKIA_GEOCODER = 'nokia'
+
     def setUp(self):
       self.fake_redis_connection = MockRedis()
-      self.us = user_service.UserService('user_id', redis_connection = self.fake_redis_connection)
-
-    def test_user_quota_should_be_10(self):
-      self.fake_redis_connection.hset('geocoder:user_id','geocoding_quota', 10)
-      assert self.us.user_quota() == 10
-
-    def test_should_return_0_if_negative_quota(self):
-      self.fake_redis_connection.hset('geocoder:user_id','geocoding_quota', -10)
-      assert self.us.user_quota() == 0
-
-    def test_should_return_0_if_not_user(self):
-      assert self.us.user_quota() == 0
 
     def test_user_used_quota_for_a_month(self):
-      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.us.used_quota_month(2015, 11) == 20
+      us = self.__build_user_service()
+      self.__increment_month_geocoder_uses('test_user', '20151111')
+      assert us.used_quota(self.NOKIA_GEOCODER, 2015, 11) == 20
+
+    def test_org_used_quota_for_a_month(self):
+      us = self.__build_user_service(organization=True)
+      self.__increment_month_geocoder_uses('test_user', '20151111', org=True)
+      assert us.used_quota(self.NOKIA_GEOCODER, 2015, 11) == 20
 
     def test_user_not_amount_in_used_quota_for_month_should_be_0(self):
-      assert self.us.used_quota_month(2015, 11) == 0
+      us = self.__build_user_service()
+      assert us.used_quota(self.NOKIA_GEOCODER, 2015, 11) == 0
 
-    def test_increment_used_quota(self):
-      self.us.increment_geocoder_use(2015, 11, 'tx_id', 1)
-      assert self.us.used_quota_month(2015, 11) == 1
+    def test_org_not_amount_in_used_quota_for_month_should_be_0(self):
+      us = self.__build_user_service(organization=True)
+      assert us.used_quota(self.NOKIA_GEOCODER, 2015, 11) == 0
+
+    def test_should_increment_user_used_quota(self):
+      us = self.__build_user_service()
+      date = datetime.strptime("20151111", "%Y%m%d")
+      us.increment_service_use(self.NOKIA_GEOCODER, date=date, amount=1)
+      assert us.used_quota(self.NOKIA_GEOCODER, 2015, 11) == 1
+      assert us.used_quota(self.NOKIA_GEOCODER, 2015, 11, 11) == 1
+
+    def test_should_increment_org_used_quota(self):
+      us = self.__build_user_service(organization=True)
+      date = datetime.strptime("20151111", "%Y%m%d")
+      us.increment_service_use(self.NOKIA_GEOCODER, date=date, amount=1)
+      assert us.used_quota(self.NOKIA_GEOCODER, 2015, 11) == 1
+      assert us.used_quota(self.NOKIA_GEOCODER, 2015, 11, 11) == 1
 
     def test_exception_if_not_redis_config(self):
-      assert_raises(Exception, user_service.UserService, 'user_id')
+      assert_raises(Exception, user_service.UserService, 'user_id')
+
+    def __build_user_service(self, organization=False, service='nokia'):
+      is_organization = 'true' if organization else 'false'
+      user_config_json = '{{"is_organization": {0}, "entity_name": "test_user"}}'.format(is_organization)
+      user_config = config_helper.UserConfig(user_config_json, 'user_id')
+
+      return user_service.UserService(user_config, service, redis_connection = self.fake_redis_connection)
+
+    def __increment_month_geocoder_uses(self, entity_name, date_string, service='nokia', amount=20, org=False):
+      parent_tag = 'org' if org else 'user'
+      date = datetime.strptime(date_string, "%Y%m%d")
+      redis_name = "{0}:{1}:{2}:used_quota_month".format(parent_tag, entity_name, service)
+      redis_key_month = "{0}_{1}".format(date.year, date.month)
+      self.fake_redis_connection.hset(redis_name, redis_key_month, amount)