class ParseTest(TestCase):
def setUp(self):
self.t = Translator()
def testSuccess(self):
r = self.t.parse_reply(b'OK true 575 60\n')
self.assertIsInstance(r, Response)
self.assertTrue(r.is_allowed)
self.assertEqual(575, r.current_credit)
self.assertEqual(60, r.next_reset_seconds)
def testFailure(self):
r = self.t.parse_reply(b'OK false 0 13\n')
self.assertIsInstance(r, Response)
self.assertFalse(r.is_allowed)
self.assertEqual(0, r.current_credit)
self.assertEqual(13, r.next_reset_seconds)
def testServerError(self):
try:
reply = b'ERR unknown-command "Unrecognized command: GETBACK"\n'
r = self.t.parse_reply(reply)
self.fail("parse_reply should have thrown a ServerError")
except ServerError as e:
self.assertEqual("unknown-command", e.error_code)
self.assertEqual("Unrecognized command: GETBACK", e.message)
def testInvalidMessage(self):
self.assertRaises(ParseError, self.t.parse_reply, b'foo true 550 50\n')
self.assertRaises(ParseError, self.t.parse_reply, b'OK foo 550 50\n')
self.assertRaises(ParseError, self.t.parse_reply, b'OK true foo 50\n')
self.assertRaises(ParseError, self.t.parse_reply, b'OK true 550 foo\n')
def setUp(self):
self.savedReactor = twisted_client.reactor
self.clock = task.Clock()
twisted_client.reactor = self.clock
self.transport = proto_helpers.StringTransport()
self.factory = twisted_client.DivvyFactory(timeout=30)
self.protocol = self.factory.buildProtocol(('127.0.0.1', 0))
self.protocol.makeConnection(self.transport)
self.translator = Translator()
def setUp(self):
self.host = "127.0.0.1"
self.factory = DivvyServerFactory()
self.listeningPort = reactor.listenTCP(0,
self.factory,
interface=self.host)
self.addCleanup(self.listeningPort.stopListening)
self.port = self.listeningPort.getHost().port
self.translator = Translator()
self.client = None
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 DivvyProtocolTest(unittest.TestCase):
def setUp(self):
self.savedReactor = twisted_client.reactor
self.clock = task.Clock()
twisted_client.reactor = self.clock
self.transport = proto_helpers.StringTransport()
self.factory = twisted_client.DivvyFactory(timeout=30)
self.protocol = self.factory.buildProtocol(('127.0.0.1', 0))
self.protocol.makeConnection(self.transport)
self.translator = Translator()
def tearDown(self):
twisted_client.reactor = self.savedReactor
def test_checkRateLimit(self):
# tx
d = self.factory.checkRateLimit({})
txData = self.transport.value()
self.assertEqual(txData, b'HIT\n')
# rx
self.transport.clear()
rxData = b'OK true 575 60\n'
response = self.translator.parse_reply(rxData)
d.addCallback(self.assertEqual, response)
self.protocol.dataReceived(rxData)
return d
def test_timeout(self):
d = self.factory.checkRateLimit({})
self.clock.advance(self.factory.timeout)
return self.assertFailure(d, TimeoutError)
def __init__(self, timeout=1.0, encoding='utf-8', debug_mode=False):
self.timeout = timeout
self.translator = Translator(encoding)
self.deferred = Deferred()
self.deferredResponses = deque()
self.divvyProtocol = None
self.running = True
self.addr = None
self.debug_mode = debug_mode
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
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")
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 __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,
socket_read_size=1024, encoding='utf-8'):
self.host = host
self.port = port
self.socket_timeout = socket_timeout
self.socket_connect_timeout = socket_connect_timeout
self.socket_keepalive = socket_keepalive
self.socket_keepalive_options = socket_keepalive_options or {}
self.socket_type = socket_type
self.retry_on_timeout = retry_on_timeout
self.socket_read_size = socket_read_size
self._translator = Translator(encoding)
self._sock = None
class RemoteDivvyProtocolTest(unittest.TestCase):
def setUp(self):
self.host = "127.0.0.1"
self.factory = DivvyServerFactory()
self.listeningPort = reactor.listenTCP(0,
self.factory,
interface=self.host)
self.addCleanup(self.listeningPort.stopListening)
self.port = self.listeningPort.getHost().port
self.translator = Translator()
self.client = None
def tearDown(self):
if self.client:
self.client.disconnect()
self.listeningPort.stopListening
reactor.disconnectAll()
for p in reactor.getDelayedCalls():
p.cancel()
def _getClientConnection(self):
"""Get a new client connection
"""
self.client = twisted_client.DivvyClient(self.host,
self.port,
timeout=1.0)
return self.client.connection.deferred
def _addRequest(self, d):
"""Make a request and check if the response is correct
"""
response = self.translator.parse_reply(self.factory.result)
d.addCallback(lambda _: self.client.check_rate_limit())
d.addCallback(self.assertEqual, response)
return d
def _pause(self, _, duration):
return task.deferLater(reactor, duration, lambda: None)
def testConnectionMade(self):
"""Validates the setup only
"""
d = self._getClientConnection()
d.addCallback(lambda _: self.client.disconnect())
return d
def testSimpleRequest(self):
"""Connect, send request and check response
"""
d = self._getClientConnection()
self._addRequest(d)
return d
def testMultipleRequests(self):
"""Send multiple request as fast as possible and check results
"""
response = self.translator.parse_reply(self.factory.result)
d = self._getClientConnection()
for _ in range(1000):
d.addCallback(lambda _: self.client.check_rate_limit())
d.addCallback(self.assertEqual, response)
return d
def testConnectionLost(self):
"""Check immediate failure if connection is down
"""
d = self.testSimpleRequest()
d.addCallback(lambda _: self.factory.crash())
d.addCallback(lambda _: self.client.check_rate_limit())
d.addTimeout(0.01, reactor)
self.assertFailure(d, ConnectionClosed)
return d
def testReconnectionOnConnectionLost(self):
"""Verify if reconnect and get a sucessful reponse after resuming
a connection lost
"""
d = self.testConnectionLost()
d.addCallback(lambda _: self.factory.resume())
d.addCallback(lambda _: self.client.connection.deferred)
self._addRequest(d)
return d
def testServerShutdown(self):
"""Check immediate failure on server shutdown
"""
d = self.testSimpleRequest()
d.addCallback(lambda _: self.factory.shutdown())
d.addCallback(lambda _: self.listeningPort.stopListening())
d.addCallback(lambda _: self.client.check_rate_limit())
d.addTimeout(0.01, reactor)
return self.assertFailure(d, ConnectionClosed)
def testReconnectionOnServerResume(self):
"""Check reconnection and sucessful request after resuming server
"""
def resume(_):
self.listeningPort = reactor.listenTCP(self.port,
self.factory,
interface=self.host)
self.factory.resume()
return self.client.connection.deferred
d = self.testServerShutdown()
d.addCallback(resume)
self._addRequest(d)
return d
def testSecondRequestAfterReconnection(self):
d = self.testReconnectionOnConnectionLost()
def cb(_):
d.addTimeout(0.01, reactor)
for _ in range(3):
d.addCallback(cb)
self._addRequest(d)
return d
def __init__(self, timeout=1.0, encoding='utf-8'):
super(DivvyProtocol, self).__init__()
self.timeout = timeout
self.translator = Translator(encoding)
self.deferred = None
def setUp(self):
self.t = Translator()
from twisted.internet import reactor
from twisted.internet.defer import Deferred, succeed, fail
from twisted.internet.protocol import ReconnectingClientFactory
from twisted.internet.endpoints import TCP4ClientEndpoint, connectProtocol
from twisted.internet.error import TimeoutError, ConnectionDone, ConnectionLost
from twisted.internet.task import deferLater
from twisted.logger import Logger
from twisted.python import log
from twisted.python.failure import Failure
from twisted.protocols.basic import LineOnlyReceiver
from twisted.protocols.policies import TimeoutMixin
from divvy.protocol import Translator
translator = Translator()
class DivvyClient(object):
log = Logger(__name__)
def __init__(self, host, port, timeout=1.0, encoding='utf-8', debug_mode=False):
"""
Configures a client that can speak to a Divvy rate limiting server.
"""
self.host = host
self.port = port
self.timeout = timeout
self.encoding = encoding
self.connection = None
self.debug_mode = debug_mode
# this may be a unwanted side effect for a constructor