def _limit_class(config, tenant, klass=None):
"""
Set up or query limit classes associated with tenants.
:param config: Name of the configuration file, for connecting to
the Redis database.
:param tenant: The ID of the tenant.
:param klass: If provided, the name of the class to map the tenant
to.
Returns the class associated with the given tenant.
"""
# Connect to the database...
db, _limits_key, _control_channel = tools.parse_config(config)
# Get the key for the limit class...
key = 'limit-class:%s' % tenant
# Now, look up the tenant's current class
old_klass = db.get(key) or 'default'
# Do we need to change it?
if klass and klass != old_klass:
if klass == 'default':
# Resetting to the default
db.delete(key)
else:
# Changing to a new value
db.set(key, klass)
return old_klass
def _group_class(config, group, klass=None, delete=False):
"""
Set up or query limit classes associated with groups.
:param config: Name of the configuration file, for connecting to
the Redis database.
:param group: The name of the group.
:param klass: If provided, the name of the class to map the group
to.
:param delete: If True, deletes the group from the database.
Returns the class associated with the given group. Note that only
one of `klass` or `delete` may be given.
"""
# Connect to the database...
db, _limits_key, _control_channel = tools.parse_config(config)
# Get the key for the limit class...
key = 'rs-group:%s' % group
# Now, look up the tenant's current class
old_klass = db.get(key)
# Do we need to delete it? Change it?
if delete:
db.delete(key)
elif klass and klass != old_klass:
db.set(key, klass)
return old_klass
def test_missing_connection(self):
with self.assertRaises(Exception):
result = tools.parse_config('bad_config')
def test_alt_control(self):
result = tools.parse_config('alt_control')
self.assertEqual(result,
(dict(host='example.com'), 'limits', 'alternate'))
def test_basic(self):
result = tools.parse_config('good_config')
self.assertEqual(result,
(dict(host='example.com'), 'limits', 'control'))
