def test_close_disconnects(self):
c = BrokerConnection(transport=Transport).channel()
conn1 = c.client.connection
conn2 = c.subclient.connection
c.close()
self.assertTrue(conn1.disconnected)
self.assertTrue(conn2.disconnected)
def test_close_disconnects(self):
c = BrokerConnection(transport=Transport).channel()
conn1 = c.client.connection
conn2 = c.subclient.connection
c.close()
self.assertTrue(conn1.disconnected)
self.assertTrue(conn2.disconnected)
def test_db_port(self):
c1 = BrokerConnection(port=None, transport=Transport).channel()
self.assertEqual(c1.client.port, Transport.default_port)
c1.close()
c2 = BrokerConnection(port=9999, transport=Transport).channel()
self.assertEqual(c2.client.port, 9999)
c2.close()
def test_db_port(self):
c1 = BrokerConnection(port=None, transport=Transport).channel()
self.assertEqual(c1.client.port, Transport.default_port)
c1.close()
c2 = BrokerConnection(port=9999, transport=Transport).channel()
self.assertEqual(c2.client.port, 9999)
c2.close()
def test__enter____exit__(self):
conn = BrokerConnection(transport=Transport)
context = conn.__enter__()
self.assertIs(context, conn)
conn.connect()
self.assertTrue(conn.connection.connected)
conn.__exit__()
self.assertIsNone(conn.connection)
conn.close() # again
def test__enter____exit__(self):
conn = BrokerConnection(backend_cls=Backend)
context = conn.__enter__()
self.assertIs(context, conn)
conn.connect()
self.assertTrue(conn.connection.connected)
conn.__exit__()
self.assertIsNone(conn.connection)
conn.close() # again
def test_establish_connection(self):
conn = BrokerConnection(port=5672, backend_cls=Backend)
conn.connect()
self.assertTrue(conn.connection.connected)
self.assertEqual(conn.host, "localhost:5672")
channel = conn.channel()
self.assertTrue(channel.open)
self.assertEqual(conn.drain_events(), "event")
_connection = conn.connection
conn.close()
self.assertFalse(_connection.connected)
self.assertIsInstance(conn.backend, Backend)
def test_close_survives_connerror(self):
class _CustomError(Exception):
pass
class MyTransport(Transport):
connection_errors = (_CustomError, )
def close_connection(self, connection):
raise _CustomError("foo")
conn = BrokerConnection(transport=MyTransport)
conn.connect()
conn.close()
self.assertTrue(conn._closed)
def test_close_survives_connerror(self):
class _CustomError(Exception):
pass
class MyTransport(Transport):
connection_errors = (_CustomError, )
def close_connection(self, connection):
raise _CustomError("foo")
conn = BrokerConnection(transport=MyTransport)
conn.connect()
conn.close()
self.assertTrue(conn._closed)
class AMQPWorker(Worker):
queues = [
{'routing_key': 'test',
'name': 'test',
'handler': 'handle_test'
}
]
_connection = None
def handle_test(self, body, message):
log.debug("Handle message: %s" % body)
message.ack()
def handle(self):
log.debug("Start consuming")
exchange = Exchange('amqp.topic', type='direct', durable=True)
self._connection = BrokerConnection(*CONNECTION)
channel = self._connection.channel()
for entry in self.queues:
log.debug("prepare to consume %s" % entry['routing_key'])
queue = Queue(entry['name'], exchange=exchange,
routing_key=entry['routing_key'])
consumer = Consumer(channel, queue)
consumer.register_callback(getattr(self, entry['handler']))
consumer.consume()
log.debug("start consuming...")
while True:
try:
self._connection.drain_events()
except socket.timeout:
log.debug("nothing to consume...")
break
self._connection.close()
def run(self):
while self.alive:
try:
self.handle()
except Exception:
self.alive = False
raise
def handle_quit(self, sig, frame):
if self._connection is not None:
self._connection.close()
self.alive = False
def handle_exit(self, sig, frame):
if self._connection is not None:
self._connection.close()
self.alive = False
sys.exit(0)
class AMQPWorker(Worker):
queues = [{
'routing_key': 'test',
'name': 'test',
'handler': 'handle_test'
}]
_connection = None
def handle_test(self, body, message):
log.debug("Handle message: %s" % body)
message.ack()
def handle(self):
log.debug("Start consuming")
exchange = Exchange('amqp.topic', type='direct', durable=True)
self._connection = BrokerConnection(*CONNECTION)
channel = self._connection.channel()
for entry in self.queues:
log.debug("prepare to consume %s" % entry['routing_key'])
queue = Queue(entry['name'],
exchange=exchange,
routing_key=entry['routing_key'])
consumer = Consumer(channel, queue)
consumer.register_callback(getattr(self, entry['handler']))
consumer.consume()
log.debug("start consuming...")
while True:
try:
self._connection.drain_events()
except socket.timeout:
log.debug("nothing to consume...")
break
self._connection.close()
def run(self):
while self.alive:
try:
self.handle()
except Exception:
self.alive = False
raise
def handle_quit(self, sig, frame):
if self._connection is not None:
self._connection.close()
self.alive = False
def handle_exit(self, sig, frame):
if self._connection is not None:
self._connection.close()
self.alive = False
sys.exit(0)
def test_close_poller_not_active(self):
c = BrokerConnection(transport=Transport).channel()
c.client.connection
c.close()
self.assertFalse(c._poller.isAlive())
self.assertTrue("BGSAVE" in c.client)
def test_close_ResponseError(self):
c = BrokerConnection(transport=Transport).channel()
c.client.bgsave_raises_ResponseError = True
c.close()
def test_close_poller_not_active(self):
c = BrokerConnection(transport=Transport).channel()
cycle = c.connection.cycle
c.client.connection
c.close()
self.assertNotIn(c, cycle._channels)
class test_Redis(TestCase):
def setUp(self):
self.connection = BrokerConnection(transport=Transport)
self.exchange = Exchange("test_Redis", type="direct")
self.queue = Queue("test_Redis", self.exchange, "test_Redis")
def tearDown(self):
self.connection.close()
def test_publish__get(self):
channel = self.connection.channel()
producer = Producer(channel, self.exchange, routing_key="test_Redis")
self.queue(channel).declare()
producer.publish({"hello": "world"})
self.assertDictEqual(
self.queue(channel).get().payload, {"hello": "world"})
self.assertIsNone(self.queue(channel).get())
self.assertIsNone(self.queue(channel).get())
self.assertIsNone(self.queue(channel).get())
def test_publish__consume(self):
connection = BrokerConnection(transport=Transport)
channel = connection.channel()
producer = Producer(channel, self.exchange, routing_key="test_Redis")
consumer = Consumer(channel, self.queue)
producer.publish({"hello2": "world2"})
_received = []
def callback(message_data, message):
_received.append(message_data)
message.ack()
consumer.register_callback(callback)
consumer.consume()
self.assertIn(channel, channel.connection.cycle._channels)
try:
connection.drain_events(timeout=1)
self.assertTrue(_received)
with self.assertRaises(socket.timeout):
connection.drain_events(timeout=0.01)
finally:
channel.close()
def test_purge(self):
channel = self.connection.channel()
producer = Producer(channel, self.exchange, routing_key="test_Redis")
self.queue(channel).declare()
for i in range(10):
producer.publish({"hello": "world-%s" % (i, )})
self.assertEqual(channel._size("test_Redis"), 10)
self.assertEqual(self.queue(channel).purge(), 10)
channel.close()
def test_db_values(self):
c1 = BrokerConnection(virtual_host=1, transport=Transport).channel()
self.assertEqual(c1.client.db, 1)
c2 = BrokerConnection(virtual_host="1", transport=Transport).channel()
self.assertEqual(c2.client.db, 1)
c3 = BrokerConnection(virtual_host="/1", transport=Transport).channel()
self.assertEqual(c3.client.db, 1)
with self.assertRaises(Exception):
BrokerConnection(virtual_host="/foo",
transport=Transport).channel()
def test_db_port(self):
c1 = BrokerConnection(port=None, transport=Transport).channel()
self.assertEqual(c1.client.port, Transport.default_port)
c1.close()
c2 = BrokerConnection(port=9999, transport=Transport).channel()
self.assertEqual(c2.client.port, 9999)
c2.close()
def test_close_poller_not_active(self):
c = BrokerConnection(transport=Transport).channel()
cycle = c.connection.cycle
c.client.connection
c.close()
self.assertNotIn(c, cycle._channels)
def test_close_ResponseError(self):
c = BrokerConnection(transport=Transport).channel()
c.client.bgsave_raises_ResponseError = True
c.close()
def test_close_disconnects(self):
c = BrokerConnection(transport=Transport).channel()
conn1 = c.client.connection
conn2 = c.subclient.connection
c.close()
self.assertTrue(conn1.disconnected)
self.assertTrue(conn2.disconnected)
def test_get__Empty(self):
channel = self.connection.channel()
with self.assertRaises(Empty):
channel._get("does-not-exist")
channel.close()
def test_get_client(self):
myredis, exceptions = _redis_modules()
@module_exists(myredis, exceptions)
def _do_test():
conn = BrokerConnection(transport=Transport)
chan = conn.channel()
self.assertTrue(chan.Client)
self.assertTrue(chan.ResponseError)
self.assertTrue(conn.transport.connection_errors)
self.assertTrue(conn.transport.channel_errors)
_do_test()
class test_Redis(TestCase):
def setUp(self):
self.connection = BrokerConnection(transport=Transport)
self.exchange = Exchange("test_Redis", type="direct")
self.queue = Queue("test_Redis", self.exchange, "test_Redis")
def tearDown(self):
self.connection.close()
def test_publish__get(self):
channel = self.connection.channel()
producer = Producer(channel, self.exchange, routing_key="test_Redis")
self.queue(channel).declare()
producer.publish({"hello": "world"})
self.assertDictEqual(self.queue(channel).get().payload,
{"hello": "world"})
self.assertIsNone(self.queue(channel).get())
self.assertIsNone(self.queue(channel).get())
self.assertIsNone(self.queue(channel).get())
def test_publish__consume(self):
connection = BrokerConnection(transport=Transport)
channel = connection.channel()
producer = Producer(channel, self.exchange, routing_key="test_Redis")
consumer = Consumer(channel, self.queue)
producer.publish({"hello2": "world2"})
_received = []
def callback(message_data, message):
_received.append(message_data)
message.ack()
consumer.register_callback(callback)
consumer.consume()
self.assertIn(channel, channel.connection.cycle._channels)
try:
connection.drain_events(timeout=1)
self.assertTrue(_received)
with self.assertRaises(socket.timeout):
connection.drain_events(timeout=0.01)
finally:
channel.close()
def test_purge(self):
channel = self.connection.channel()
producer = Producer(channel, self.exchange, routing_key="test_Redis")
self.queue(channel).declare()
for i in range(10):
producer.publish({"hello": "world-%s" % (i, )})
self.assertEqual(channel._size("test_Redis"), 10)
self.assertEqual(self.queue(channel).purge(), 10)
channel.close()
def test_db_values(self):
c1 = BrokerConnection(virtual_host=1,
transport=Transport).channel()
self.assertEqual(c1.client.db, 1)
c2 = BrokerConnection(virtual_host="1",
transport=Transport).channel()
self.assertEqual(c2.client.db, 1)
c3 = BrokerConnection(virtual_host="/1",
transport=Transport).channel()
self.assertEqual(c3.client.db, 1)
with self.assertRaises(Exception):
BrokerConnection(virtual_host="/foo",
transport=Transport).channel()
def test_db_port(self):
c1 = BrokerConnection(port=None, transport=Transport).channel()
self.assertEqual(c1.client.port, Transport.default_port)
c1.close()
c2 = BrokerConnection(port=9999, transport=Transport).channel()
self.assertEqual(c2.client.port, 9999)
c2.close()
def test_close_poller_not_active(self):
c = BrokerConnection(transport=Transport).channel()
cycle = c.connection.cycle
c.client.connection
c.close()
self.assertNotIn(c, cycle._channels)
def test_close_ResponseError(self):
c = BrokerConnection(transport=Transport).channel()
c.client.bgsave_raises_ResponseError = True
c.close()
def test_close_disconnects(self):
c = BrokerConnection(transport=Transport).channel()
conn1 = c.client.connection
conn2 = c.subclient.connection
c.close()
self.assertTrue(conn1.disconnected)
self.assertTrue(conn2.disconnected)
def test_get__Empty(self):
channel = self.connection.channel()
with self.assertRaises(Empty):
channel._get("does-not-exist")
channel.close()
def test_get_client(self):
myredis, exceptions = _redis_modules()
@module_exists(myredis, exceptions)
def _do_test():
conn = BrokerConnection(transport=Transport)
chan = conn.channel()
self.assertTrue(chan.Client)
self.assertTrue(chan.ResponseError)
self.assertTrue(conn.transport.connection_errors)
self.assertTrue(conn.transport.channel_errors)
_do_test()
def test_close_poller_not_active(self):
c = BrokerConnection(transport=Transport).channel()
cycle = c.connection.cycle
c.client.connection
c.close()
self.assertNotIn(c, cycle._channels)
def test_close_poller_not_active(self):
c = BrokerConnection(transport=Transport).channel()
c.client.connection
c.close()
self.assertFalse(c._poller.isAlive())
self.assertTrue("BGSAVE" in c.client)
def test_close_ResponseError(self):
c = BrokerConnection(transport=Transport).channel()
c.client.bgsave_raises_ResponseError = True
c.close()
class Connection:
def __init__(self, hostname, port, vhost, userid, password):
"""connects to broker and provides convenience methods"""
self.broker = BrokerConnection(hostname=hostname, port=port,
userid=userid, password=password,
virtual_host=vhost)
def __del__(self):
self.broker.close()
def declare(self, exchange, exchange_type, binding="", queue=""):
"""declares the exchange, the queue and binds the queue to the exchange
exchange - exchange name
exchange_type - direct, topic, fanout
binding - binding to queue (optional)
queue - queue to bind to exchange using binding (optional)
"""
if (binding and not queue) or (queue and not binding):
if queue and not exchange_type == "fanout":
raise Error("binding and queue are not mutually exclusive")
consumer = Consumer(connection=self.broker,
exchange=exchange, exchange_type=exchange_type,
routing_key=binding, queue=queue)
consumer.declare()
consumer.close()
def consume(self, queue, limit=None, callback=None, auto_declare=False):
"""consume messages in queue
queue - name of queue
limit - amount of messages to iterate through (default: no limit)
callback - method to call when a new message is received
must take two arguments: message_data, message
must send the acknowledgement: message.ack()
default: print message to stdout and send ack
auto_declare - automatically declare the queue (default: false)
"""
if not callback:
callback = _consume_callback
consumer = Consumer(connection=self.broker, queue=queue,
auto_declare=auto_declare)
consumer.register_callback(callback)
for message in consumer.iterqueue(limit=limit, infinite=False):
consumer.receive(message.payload, message)
consumer.close()
def publish(self, exchange, routing_key, message,
auto_declare=False, persistent=True):
"""publish a message to exchange using routing_key
exchange - name of exchange
routing_key - interpretation of routing key depends on exchange type
message - message content to send
auto_declare - automatically declare the exchange (default: false)
persistent - store message on disk as well as memory (default: True)
"""
delivery_mode = 2
if not persistent:
delivery_mode = 1
publisher = Publisher(connection=self.broker,
exchange=exchange, routing_key=routing_key,
auto_declare=auto_declare)
publisher.send(message, delivery_mode=delivery_mode)
publisher.close()
password='******',
virtual_host='/'
)
channel = connection.channel()
consumer = Consumer(channel, server_queue)
consumer.register_callback(receive_message)
consumer.consume(no_ack=True)
while True:
try:
connection.drain_events(timeout=CONSUME_TIMEOUT)
except socket.timeout:
# flush the server buffer if we haven't gotten
# any messages in a while
if len(server_buffer) > 0:
logging.debug('QUEUE TIMED OUT, FLUSHING BUFFER')
process_servers()
elif display_waiting:
processing_time = datetime.now() -\
last_waiting_time -\
timedelta(seconds=CONSUME_TIMEOUT)
logging.debug(
'... WAITING ... TIME SINCE LAST: %s',
processing_time
)
display_waiting = False
connection.close()
