class BuildTest(TestCase):
def setUp(self):
self.t = Translator()
def test_no_arguments(self):
output = self.t.build_hit()
self.assertEqual(b'HIT\n', output)
def test_single_argument(self):
output = self.t.build_hit(method="GET")
self.assertEqual(b'HIT "method"="GET"\n', output)
def test_multiple_arguments(self):
output = self.t.build_hit(path="/cookies", method="GET")
self.assertEqual(b'HIT "method"="GET" "path"="/cookies"\n', output)
def test_coerced_arguments(self):
output = self.t.build_hit(b=True, f=3.1416, i=-65535)
self.assertEqual(b'HIT "b"="True" "f"="3.1416" "i"="-65535"\n', output)
def test_special_keys(self):
kwargs = {'Arg+1': 1, 'arg-*': '*'}
output = self.t.build_hit(**kwargs)
self.assertEqual(b'HIT "Arg+1"="1" "arg-*"="*"\n', output)
def test_encoding(self):
output = self.t.build_hit(happy_face=u'\U0001f604')
self.assertEqual(b'HIT "happy_face"="\xf0\x9f\x98\x84"\n', output)
def test_invalid_key(self):
kwargs = {'Arg 3': 3}
self.assertRaises(InputError, self.t.build_hit, **kwargs)
def test_invalid_value(self):
self.assertRaises(InputError, self.t.build_hit, method="GET PUT")
class DivvyClient(object):
def __init__(self,
host='localhost',
port=8321,
socket_timeout=1,
socket_connect_timeout=1,
socket_keepalive=False,
socket_keepalive_options=None,
socket_type=0,
retry_on_timeout=False,
encoding='utf-8'):
self.host = host
self.port = port
self.translator = Translator(encoding=encoding)
self.connection = Connection(
host=host,
port=port,
socket_timeout=socket_timeout,
socket_connect_timeout=socket_connect_timeout,
socket_keepalive=socket_keepalive,
socket_keepalive_options=socket_keepalive_options,
socket_type=socket_type,
retry_on_timeout=retry_on_timeout)
def check_rate_limit(self, **kwargs):
"""Perform a check-and-decrement of quota. Zero or more key-value pairs
specify the operation being performed, and will be evaluated by the
server against its configuration.
Args:
**kwargs: Zero or more key-value pairs to specify the operation
being performed, which will be evaluated by the server against
its configuration.
Returns:
divvy.Response, a namedtuple with these fields:
is_allowed: one of true, false, indicating whether quota was
available.
current_credit: number of credit(s) available at the end of
this command.
next_reset_seconds: time, in seconds, until credit next resets.
"""
cmd = self.translator.build_hit(**kwargs)
self.connection.send(cmd)
reply = self.connection.recv()
response = self.translator.parse_reply(reply)
return response
class DivvyProtocol(LineOnlyReceiver, TimeoutMixin, object):
"""
Twisted handler for network communication with a Divvy server.
"""
delimiter = "\n"
def __init__(self, timeout=1.0, encoding='utf-8'):
super(DivvyProtocol, self).__init__()
self.timeout = timeout
self.translator = Translator(encoding)
self.deferred = None
def timeoutConnection(self):
TimeoutMixin.timeoutConnection(self)
# self.deferred should be present whenever this method is
# called, but let's program defensively and double-check.
if self.deferred:
self.deferred.errback(TimeoutError())
def checkRateLimit(self, **kwargs):
assert self.connected
assert self.deferred is None
self.deferred = Deferred()
line = self.translator.build_hit(**kwargs).strip()
self.sendLine(line)
self.setTimeout(self.timeout)
return self.deferred
def lineReceived(self, line):
self.setTimeout(None)
# we should never receive a line if there's no outstanding request
assert self.deferred is not None
try:
response = self.translator.parse_reply(line)
self.deferred.callback(response)
except Exception as e:
self.deferred.errback(e)
finally:
self.transport.loseConnection()
self.deferred = None