├── tests ├── __init__.py └── rabbit │ ├── test_publisher.py │ ├── test_subscriber.py │ └── test_queue.py ├── equeue ├── rabbit │ ├── __init__.py │ ├── subscriber.py │ ├── publisher.py │ └── queue.py └── __init__.py ├── MANIFEST.in ├── requirements.txt ├── requirements-dev.txt ├── .gitignore ├── setup.py ├── LICENSE ├── example.py └── README.md /tests/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /equeue/rabbit/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /equeue/__init__.py: -------------------------------------------------------------------------------- 1 | __version__ = '0.1.5' 2 | -------------------------------------------------------------------------------- /MANIFEST.in: -------------------------------------------------------------------------------- 1 | include README.md 2 | include LICENSE 3 | -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | amqp==2.1.0 2 | simplejson==3.8.2 3 | -------------------------------------------------------------------------------- /requirements-dev.txt: -------------------------------------------------------------------------------- 1 | mock==2.0.0 2 | pytest==3.0.3 3 | pytest-cov==2.3.1 4 | 5 | -------------------------------------------------------------------------------- /equeue/rabbit/subscriber.py: -------------------------------------------------------------------------------- 1 | # enconding: utf-8 2 | import time 3 | from equeue.rabbit.queue import RabbitQueue 4 | 5 | 6 | class Subscriber(RabbitQueue): 7 | 8 | def get(self, queue_name=None, block=True): 9 | """ 10 | Gets messages from the queue. 11 | queue_name: optional, defaults to the one passed on __init__(). 12 | block: boolean. If block is True (default), get() will not return until 13 | a message is acquired from the queue. 14 | 15 | Returns a tuple (message, delivery_tag) when a message is read, where 16 | message is a deserialized json and delivery_tag is a parameter used 17 | for on ack() and reject() methods. If block is False and there's no 18 | message on the queue, returns (None, None). 19 | """ 20 | while True: 21 | message = self._try('basic_get', queue=queue_name or self.default_queue_name) 22 | if message: 23 | parsed = self._parse_message(message) 24 | if parsed: 25 | return parsed 26 | 27 | if block: 28 | time.sleep(.5) 29 | else: 30 | return None, None 31 | 32 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # Byte-compiled / optimized / DLL files 2 | __pycache__/ 3 | *.py[cod] 4 | 5 | # C extensions 6 | *.so 7 | 8 | # Distribution / packaging 9 | .Python 10 | env/ 11 | build/ 12 | develop-eggs/ 13 | dist/ 14 | downloads/ 15 | eggs/ 16 | .eggs/ 17 | lib/ 18 | lib64/ 19 | parts/ 20 | sdist/ 21 | var/ 22 | *.egg-info/ 23 | .installed.cfg 24 | *.egg 25 | 26 | # PyInstaller 27 | # Usually these files are written by a python script from a template 28 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 29 | *.manifest 30 | *.spec 31 | 32 | # Installer logs 33 | pip-log.txt 34 | pip-delete-this-directory.txt 35 | 36 | # Unit test / coverage reports 37 | htmlcov/ 38 | .tox/ 39 | .coverage 40 | .coverage.* 41 | .cache 42 | nosetests.xml 43 | coverage.xml 44 | *,cover 45 | 46 | # Translations 47 | *.mo 48 | *.pot 49 | 50 | # Django stuff: 51 | *.log 52 | 53 | # Sphinx documentation 54 | docs/_build/ 55 | 56 | # PyBuilder 57 | target/ 58 | *.py[co] 59 | 60 | *.swp 61 | 62 | #Backups 63 | *~ 64 | 65 | # Packages 66 | *.egg 67 | *.egg-info 68 | dist 69 | build 70 | eggs 71 | parts 72 | bin 73 | develop-eggs 74 | .installed.cfg 75 | 76 | # Installer logs 77 | pip-log.txt 78 | 79 | # Unit test / coverage reports 80 | .coverage 81 | .tox 82 | 83 | db.sqlite3 84 | static/collect/* 85 | configmanager/settings_local.py 86 | 87 | cache/ 88 | output/ 89 | *.pid 90 | .idea/ 91 | .DS_Store 92 | .vagrant/ 93 | *.retry -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | # enconding: utf-8 3 | import sys 4 | import os 5 | from setuptools import setup, find_packages 6 | 7 | from equeue import __version__ 8 | 9 | BASE_PATH = os.path.dirname(__file__) 10 | setup( 11 | name='equeue', 12 | version=__version__, 13 | description='Elastic Queue is a library to bear a PubSub projects. ', 14 | long_description=open(os.path.join(BASE_PATH, 'README.md')).read(), 15 | author='Jesue Junior', 16 | author_email='jesuesousa@gmail.com', 17 | url='https://github.com/jesuejunior/equeue', 18 | packages=find_packages(), 19 | install_requires=['amqp==2.1.0', 'simplejson>=3.8.2', 'six==1.10.0'], 20 | test_suite='tests', 21 | tests_require=['tox>=2.3.1', 'pytest==3.0.3', 'pytest-cov==2.3.1'] + ( 22 | ['mock==2.0.0'] if sys.version_info.major == 2 else [] 23 | ), 24 | classifiers=[ 25 | 'Intended Audience :: Developers', 26 | 'License :: OSI Approved :: BSD License', 27 | 'Programming Language :: Python', 28 | 'Programming Language :: Python :: 2.7', 29 | 'Programming Language :: Python :: 3', 30 | 'Programming Language :: Python :: 3.5', 31 | 'Programming Language :: Python :: Implementation :: PyPy', 32 | 'Topic :: Software Development', 33 | 'Topic :: Software Development :: Libraries', 34 | 'Topic :: Software Development :: Libraries :: Python Modules', 35 | ], 36 | ) 37 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | The BSD 3-Clause License 2 | 3 | Copyright (c) 2016, Jesué Junior 4 | 5 | All rights reserved. 6 | 7 | Redistribution and use in source and binary forms, with or without modification, 8 | are permitted provided that the following conditions are met: 9 | 10 | 1. Redistributions of source code must retain the above copyright notice, 11 | this list of conditions and the following disclaimer. 12 | 13 | 2. Redistributions in binary form must reproduce the above copyright notice, 14 | this list of conditions and the following disclaimer in the documentation 15 | and/or other materials provided with the distribution. 16 | 17 | 3. Neither the name of the copyright holder nor the names of its contributors 18 | may be used to endorse or promote products derived from this software without 19 | specific prior written permission. 20 | 21 | THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" 22 | AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, 23 | THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. 24 | IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, 25 | INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, 26 | BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, 27 | DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, 28 | WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) 29 | ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 30 | -------------------------------------------------------------------------------- /example.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | import threading, logging, time 3 | import json 4 | from equeue.rabbit.publisher import Publisher 5 | from equeue.rabbit.subscriber import Subscriber 6 | 7 | def events_out(callback, message, delivery_tag): 8 | print('This a full message', message) 9 | print('This is a deliver tag', delivery_tag) 10 | callback.ack(delivery_tag) 11 | 12 | 13 | class Producer(threading.Thread): 14 | # daemon = True 15 | 16 | def run(self): 17 | pub = Publisher(host='localhost', username='guest', 18 | password='guest', queue_name='teste') 19 | while True: 20 | for i in xrange(10000): 21 | logging.info('producer') 22 | pub.put(body=json.dumps({'id': i}), routing_key='t') 23 | time.sleep(60) 24 | 25 | 26 | class Consumer(threading.Thread): 27 | # daemon = True 28 | 29 | def run(self): 30 | sub = Subscriber(host='localhost', username='everest', 31 | password='everest', queue_name='t') 32 | 33 | sub.setup_consumer(callback=events_out) 34 | 35 | while True: 36 | # time.sleep(0.01) 37 | logging.info('consumer') 38 | sub.consume() 39 | 40 | def main(): 41 | threads = [ 42 | Consumer(), 43 | Producer() 44 | ] 45 | 46 | for t in threads: 47 | t.start() 48 | # time.sleep(10) 49 | 50 | if __name__ == "__main__": 51 | logging.basicConfig( 52 | format='%(asctime)s.%(msecs)s:%(name)s:%(thread)d:%(levelname)s:%(process)d:%(message)s', 53 | level=logging.INFO 54 | ) 55 | main() 56 | -------------------------------------------------------------------------------- /equeue/rabbit/publisher.py: -------------------------------------------------------------------------------- 1 | # encondign: utf-8 2 | import simplejson as json 3 | from amqp import Message 4 | from equeue.rabbit.queue import RabbitQueue, SerializationError 5 | 6 | 7 | class Publisher(RabbitQueue): 8 | 9 | def put(self, message_dict=None, routing_key='', exchange=None, body=None, priority=0): 10 | """ 11 | Publishes a message to the queue. 12 | message_dict: the json-serializable object that will be published 13 | when body is None 14 | routing_key: the routing key for the message. 15 | exchange: the exchange to which the message will be published. 16 | body: The message to be published. If none, message_dict is published. 17 | 18 | It also works as a context manager: 19 | with Publisher(**options) as queue: 20 | for msg in msgs: 21 | queue.put(msg) 22 | """ 23 | if exchange is None: 24 | exchange = self.default_exchange or '' 25 | if body is None: 26 | try: 27 | body = json.dumps(message_dict) 28 | except Exception as e: 29 | raise SerializationError(e) 30 | 31 | message = Message(body, 32 | delivery_mode=2, 33 | content_type='application/json', 34 | priority=priority 35 | ) 36 | return self._try('basic_publish', 37 | msg=message, 38 | exchange=exchange, 39 | routing_key=routing_key) 40 | 41 | def flush(self): 42 | self.put("flush", routing_key="/dev/null") 43 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # EQueue 2 | Elastic Queue is a library to bear a PubSub projects 3 | 4 | ### How to use it 5 | 6 | Install via PIP 7 | 8 | ```shell 9 | $ pip install equeue 10 | ``` 11 | 12 | #### Producing 13 | 14 | Using ipython 15 | 16 | ```python 17 | 18 | In [1]: from equeue.rabbit.publisher import Publisher 19 | 20 | In [2]: pub = Publisher(host='localhost', username='guest', password='guest', queue_name='t') 21 | 22 | In [3]: pub.put(message_dict={'id': 1}) 23 | Out[3]: 24 | 25 | ``` 26 | 27 | #### By poll 28 | 29 | Using ipython 30 | 31 | ```python 32 | 33 | In [1]: from equeue.rabbit.subscriber import Subscriber 34 | 35 | In [2]: sub = Subscriber(host='localhost', username='guest', password='guest', queue_name='t') 36 | 37 | In [3]: msg = sub.get() 38 | 39 | In [4]: if msg: 40 | ...: print(msg) 41 | ...: 42 | ``` 43 | 44 | #### Consuming 45 | 46 | Creating a main.py you'd see better. 47 | 48 | ```python 49 | 50 | from equeue.rabbit.subscriber import Subscriber 51 | 52 | def events_out(callback, message, delivery_tag): 53 | print(message) 54 | print(delivery_tag) 55 | callback.ack(delivery_tag) 56 | 57 | def main(): 58 | 59 | sub = Subscriber(host='localhost', username='guest', 60 | password='guest', queue_name='t') 61 | 62 | sub.setup_consumer(callback=events_out) 63 | while True: 64 | sub.consume() 65 | if __name__ == '__main__': 66 | main() 67 | 68 | ``` 69 | 70 | Then 71 | 72 | ```shell 73 | $ python main.py 74 | ``` 75 | 76 | ### Developing mode 77 | 78 | Running tests 79 | 80 | To run the project's test you will need to have pytest installed. The instalation is simple as : 81 | 82 | ```shell 83 | $ pip install pytest 84 | ``` 85 | 86 | And to run the tests you need to ajust yout `PYTHONPATH` 87 | 88 | ```shell 89 | $ PYTHONPATH=equeue py.test 90 | ``` 91 | 92 | -------------------------------------------------------------------------------- /tests/rabbit/test_publisher.py: -------------------------------------------------------------------------------- 1 | #encoding: utf-8 2 | import unittest 3 | from datetime import datetime 4 | 5 | import simplejson as json 6 | from amqp import Message, AMQPError, ConnectionError 7 | from mock import MagicMock, patch, call, Mock, ANY 8 | 9 | from equeue.rabbit.queue import SerializationError 10 | from equeue.rabbit.publisher import Publisher 11 | 12 | 13 | class RabbitQueueTest(unittest.TestCase): 14 | def setUp(self): 15 | _delivery_info = {'delivery_tag': 'delivery_tag'} 16 | 17 | self.message = MagicMock() 18 | self.message.body = json.dumps({'id': 123}) 19 | self.message.delivery_info = _delivery_info 20 | 21 | self.channel_mock = MagicMock() 22 | self.channel_mock.basic_get.return_value = self.message 23 | 24 | self.connection = MagicMock() 25 | self.connection.channel.return_value = self.channel_mock 26 | 27 | self.connection_cls_patcher = patch('amqp.Connection', 28 | return_value=self.connection) 29 | self.connection_cls_mock = self.connection_cls_patcher.start() 30 | 31 | self.publisher = Publisher('localhost', 'username', 'pwd', 32 | exchange='default_exchange', 33 | queue_name='default_queue') 34 | 35 | def tearDown(self): 36 | self.connection_cls_patcher.stop() 37 | 38 | def test_put_default_exchange_if_not_supplied(self): 39 | amqp_msg = Message('body', delivery_mode=2, content_type='application/json', priority=0) 40 | self.publisher.put(body='body') 41 | self.assertEqual(self.channel_mock.basic_publish.call_args_list, 42 | [call(msg=amqp_msg, 43 | exchange='default_exchange', 44 | routing_key='')]) 45 | 46 | def test_put_serializes_message_if_necessary(self): 47 | message = {'key': 'value'} 48 | with patch('equeue.rabbit.publisher.Publisher') as MockPublisher: 49 | self.publisher.put(message_dict=message) 50 | 51 | self.assertEqual(MockPublisher.call_args_list, 52 | [call(json.dumps(message), delivery_mode=2, 53 | content_type='application/json', priority=0)]) 54 | 55 | def test_put_raises_serialization_error_if_cant_be_serialized_to_json(self): 56 | self.assertRaises(SerializationError, self.publisher.put, message_dict=ValueError) 57 | 58 | 59 | -------------------------------------------------------------------------------- /tests/rabbit/test_subscriber.py: -------------------------------------------------------------------------------- 1 | from datetime import datetime 2 | import json 3 | import unittest 4 | 5 | from amqp import Message, AMQPError, ConnectionError 6 | from mock import MagicMock, patch, call, Mock, ANY 7 | 8 | from equeue.rabbit.subscriber import Subscriber 9 | 10 | 11 | class SubscriberTest(unittest.TestCase): 12 | def setUp(self): 13 | _delivery_info = {'delivery_tag': 'delivery_tag'} 14 | 15 | self.message = MagicMock() 16 | self.message.body = json.dumps({'id': 123}) 17 | self.message.delivery_info = _delivery_info 18 | 19 | self.channel_mock = MagicMock() 20 | self.channel_mock.basic_get.return_value = self.message 21 | 22 | self.connection = MagicMock() 23 | self.connection.channel.return_value = self.channel_mock 24 | 25 | self.connection_cls_patcher = patch('amqp.Connection', 26 | return_value=self.connection) 27 | self.connection_cls_mock = self.connection_cls_patcher.start() 28 | 29 | self.subscriber = Subscriber('localhost', 'username', 'pwd', 30 | exchange='default_exchange', 31 | queue_name='default_queue') 32 | 33 | def tearDown(self): 34 | self.connection_cls_patcher.stop() 35 | 36 | def test_get_uses_default_queue_if_not_supplied(self): 37 | self.subscriber.get() 38 | self.assertEqual(self.channel_mock.basic_get.call_args_list, 39 | [call(queue='default_queue')]) 40 | 41 | def test_get_none_if_block_is_false_and_queue_is_empty(self): 42 | self.channel_mock.basic_get.return_value = None 43 | rv = self.subscriber.get(block=False) 44 | self.assertEqual(rv, (None, None)) 45 | 46 | def test_get_sleeps_and_tries_again_until_queue_is_not_empty(self): 47 | empty_rv = None 48 | self.channel_mock.basic_get.side_effect = [empty_rv, empty_rv, self.message] 49 | with patch('time.sleep') as sleep,\ 50 | patch('equeue.rabbit.queue.RabbitQueue._parse_message') as parse_message_mock: 51 | message = self.subscriber.get(queue_name='queue_name') 52 | 53 | self.assertEqual(message, parse_message_mock.return_value) 54 | self.assertEqual(parse_message_mock.call_args_list, 55 | [call(self.message)]) 56 | self.assertEqual(self.channel_mock.basic_get.call_args_list, 57 | [call(queue='queue_name'), 58 | call(queue='queue_name'), 59 | call(queue='queue_name')]) 60 | self.assertEqual(sleep.call_count, 2) 61 | 62 | def test_get_error_if_default_queue_does_not_exist(self): 63 | self.connection_cls_mock.return_value.channel.side_effect = ConnectionError 64 | with self.assertRaises(ConnectionError): 65 | self.subscriber.get() 66 | 67 | 68 | -------------------------------------------------------------------------------- /equeue/rabbit/queue.py: -------------------------------------------------------------------------------- 1 | import time 2 | import uuid 3 | from datetime import datetime 4 | 5 | import simplejson as json 6 | import amqp 7 | from amqp import Message, AMQPError, ConnectionError as AMQPConnectionError 8 | 9 | MAX_TRIES = 3 10 | META_FIELD = "_meta" 11 | # mainly, this class I took from https://github.com/sievetech/hived/blob/master/hived/queue.py 12 | # and adapted for my necessity. 13 | 14 | 15 | class ConnectionError(AMQPConnectionError): 16 | def __str__(self): 17 | return '%s' % self.message 18 | 19 | 20 | class SerializationError(Exception): 21 | def __init__(self, exc, body=None): 22 | super(Exception, self).__init__(*exc.args) 23 | self.exc = exc 24 | self.body = body 25 | 26 | def __repr__(self): 27 | return '%s: %s' % (self.exc, repr(self.body)) 28 | 29 | 30 | class RabbitQueue(object): 31 | """ 32 | For getting messages from the queue, see get() in the Subscriber class. 33 | For publishing, see put() in the Publisher class. 34 | The connection is lazy, i.e. it only happens on the first get() / put(). 35 | """ 36 | 37 | def __init__(self, host='localhost', username='guest', password='guest', 38 | virtual_host='/', exchange=None, queue_name=None, queue_heartbeat=None): 39 | self.default_exchange = exchange 40 | self.default_queue_name = queue_name 41 | self.channel = None 42 | self.subscription = None 43 | self.connection = None 44 | 45 | self.connection_parameters = { 46 | 'host': host, 47 | 'userid': username, 48 | 'password': password, 49 | 'virtual_host': virtual_host, 50 | 'heartbeat': queue_heartbeat 51 | } 52 | 53 | def __enter__(self): 54 | return self 55 | 56 | def __exit__(self, *_): 57 | self.close() 58 | return False 59 | 60 | def _connect(self): 61 | try: 62 | self.close() 63 | except Exception: 64 | pass 65 | self.connection = amqp.Connection(**self.connection_parameters) 66 | self.connection.connect() 67 | self.channel = self.connection.channel() 68 | if self.subscription: 69 | self._subscribe() 70 | 71 | def close(self): 72 | if self.connection is not None: 73 | self.connection.close() 74 | 75 | def _try(self, method, _tries=1, **kwargs): 76 | if self.channel is None: 77 | self._connect() 78 | 79 | try: 80 | # import ipdb; ipdb.set_trace() 81 | return getattr(self.channel, method)(**kwargs) 82 | except (AMQPError, IOError) as e: 83 | if _tries < MAX_TRIES: 84 | self._connect() 85 | return self._try(method, _tries + 1, **kwargs) 86 | else: 87 | raise ConnectionError(e) 88 | 89 | def _subscribe(self): 90 | self.default_queue_name = '%s_%s' % (self.subscription, uuid.uuid4()) 91 | self.channel.queue_declare(queue=self.default_queue_name, 92 | durable=False, 93 | exclusive=True, 94 | auto_delete=True) 95 | self.channel.queue_bind(exchange='notifies', 96 | queue=self.default_queue_name, 97 | routing_key=self.subscription) 98 | 99 | def subscribe(self, routing_key): 100 | self.subscription = routing_key 101 | self._connect() 102 | 103 | def _parse_message(self, message): 104 | body = message.body 105 | delivery_tag = message.delivery_info['delivery_tag'] 106 | try: 107 | message_dict = json.loads(body) 108 | message_dict.setdefault(META_FIELD, {}) 109 | except Exception: 110 | self.ack(delivery_tag) 111 | return 112 | return message_dict, delivery_tag 113 | 114 | def setup_consumer(self, callback, queue_names=None): 115 | # import ipdb; ipdb.set_trace() 116 | def message_callback(message): 117 | parsed = self._parse_message(message) 118 | if parsed: 119 | callback(self, *parsed) 120 | 121 | self._try('basic_qos', prefetch_size=0, prefetch_count=1, a_global=False) 122 | queue_names = queue_names or [self.default_queue_name] 123 | for queue_name in queue_names: 124 | self.channel.basic_consume(queue_name, callback=message_callback) 125 | 126 | def consume(self): 127 | self.connection.drain_events() 128 | 129 | def ack(self, delivery_tag): 130 | """ 131 | Acks a message from the queue. 132 | delivery_tag: second value on the tuple returned from get(). 133 | """ 134 | try: 135 | self.channel.basic_ack(delivery_tag) 136 | except AMQPError: 137 | # There's nothing we can do, we can't ack the message in 138 | # a different channel than the one we got it from 139 | pass 140 | 141 | def reject(self, delivery_tag): 142 | """ 143 | Rejects a message from the queue, i.e. returns it to the top of the queue. 144 | delivery_tag: second value on the tuple returned from get(). 145 | """ 146 | try: 147 | self.channel.basic_reject(delivery_tag, requeue=True) 148 | except AMQPError: 149 | pass # It's out of our hands already 150 | -------------------------------------------------------------------------------- /tests/rabbit/test_queue.py: -------------------------------------------------------------------------------- 1 | #encoding: utf-8 2 | import unittest 3 | from datetime import datetime 4 | 5 | import simplejson as json 6 | from amqp import Message, AMQPError, ConnectionError 7 | from mock import MagicMock, patch, call, Mock, ANY 8 | 9 | from equeue.rabbit.queue import (RabbitQueue, MAX_TRIES, SerializationError, META_FIELD) 10 | 11 | MODULE = 'equeue.rabbit.queue.' 12 | 13 | 14 | class RabbitQueueTest(unittest.TestCase): 15 | def setUp(self): 16 | _delivery_info = {'delivery_tag': 'delivery_tag'} 17 | 18 | self.message = MagicMock() 19 | self.message.body = json.dumps({'id': 123}) 20 | self.message.delivery_info = _delivery_info 21 | 22 | self.channel_mock = MagicMock() 23 | self.channel_mock.basic_get.return_value = self.message 24 | 25 | self.connection = MagicMock() 26 | self.connection.channel.return_value = self.channel_mock 27 | 28 | self.connection_cls_patcher = patch('amqp.Connection', 29 | return_value=self.connection) 30 | self.connection_cls_mock = self.connection_cls_patcher.start() 31 | 32 | self.external_queue = RabbitQueue('localhost', 'username', 'pwd', 33 | exchange='default_exchange', 34 | queue_name='default_queue') 35 | 36 | def tearDown(self): 37 | self.connection_cls_patcher.stop() 38 | 39 | def test_connect_calls_close_before_creating_a_new_connection(self): 40 | with patch(MODULE + 'RabbitQueue.close') as close_mock: 41 | self.external_queue._connect() 42 | self.assertEqual(close_mock.call_count, 1) 43 | 44 | def test_connect_ignores_close_errors(self): 45 | with patch.object(self.external_queue, 'close', side_effect=[Exception]) as mock_close: 46 | self.external_queue._connect() 47 | self.assertRaises 48 | self.assertEqual(mock_close.call_count, 1) 49 | 50 | def test_connect_subscribes_if_subscription_is_set(self): 51 | with patch.object(self.external_queue, 'close'), \ 52 | patch.object(self.external_queue, '_subscribe') as mock_subscribe: 53 | self.external_queue.subscription = 'routing_key' 54 | self.external_queue._connect() 55 | 56 | self.assertEqual(mock_subscribe.call_count, 1) 57 | 58 | def test__try_connects_if_disconnected(self): 59 | self.channel_mock.method.return_value = 'rv' 60 | rv = self.external_queue._try('method', arg='value') 61 | 62 | self.assertEqual(self.connection_cls_mock.call_count, 1) 63 | self.assertEqual(self.channel_mock.method.call_args_list, 64 | [call(arg='value')]) 65 | self.assertEqual(rv, 'rv') 66 | 67 | def test__try_tries_up_to_max_tries(self): 68 | self.channel_mock.method.side_effect = [AMQPError, AMQPError, 'rv'] 69 | rv = self.external_queue._try('method') 70 | 71 | self.assertEqual(self.channel_mock.method.call_count, MAX_TRIES) 72 | self.assertEqual(rv, 'rv') 73 | 74 | def test__try_doesnt_try_more_than_max_tries(self): 75 | self.channel_mock.method.side_effect = [AMQPError, AMQPError, AMQPError, 'rv'] 76 | self.assertRaises(ConnectionError, self.external_queue._try, 'method') 77 | 78 | def test_parse_message_deserializes_the_message_body_and_sets_meta_field(self): 79 | message, ack = self.external_queue._parse_message(self.message) 80 | self.assertEqual(message[META_FIELD], {}) 81 | self.assertEqual(ack, 'delivery_tag') 82 | 83 | def test_malformed_message_should_ack(self): 84 | queue = self.external_queue 85 | self.message.body = "{'foo': True}" 86 | with patch.object(queue, 'ack') as mock_ack: 87 | result = queue._parse_message(self.message) 88 | self.assertIsNone(result) 89 | self.assertEqual(mock_ack.call_args_list, 90 | [call(self.message.delivery_info['delivery_tag'])]) 91 | 92 | def test_setup_consumer(self): 93 | callback = Mock() 94 | self.external_queue.connection = Mock() 95 | self.external_queue.channel = Mock() 96 | with patch(MODULE + 'RabbitQueue._try') as try_mock: 97 | self.external_queue.setup_consumer(callback, ['queue_1', 'queue_2']) 98 | 99 | self.assertEqual(try_mock.call_args_list, 100 | [call('basic_qos', prefetch_size=0, 101 | prefetch_count=1, a_global=False)]) 102 | self.assertEqual(self.external_queue.channel.basic_consume.call_args_list, 103 | [call('queue_1', callback=ANY), 104 | call('queue_2', callback=ANY)]) 105 | 106 | def test_consume(self): 107 | self.external_queue.connection = Mock() 108 | self.external_queue.consume() 109 | self.assertEqual(self.external_queue.connection.drain_events.call_count, 1) 110 | 111 | def test_ack_ignores_connection_errors(self): 112 | self.external_queue.channel = self.channel_mock 113 | self.channel_mock.basic_ack.side_effect = AMQPError 114 | self.external_queue.ack('delivery_tag') 115 | 116 | def test_reject_ignores_connection_errors(self): 117 | self.external_queue.channel = self.channel_mock 118 | self.channel_mock.basic_reject.side_effect = AMQPError 119 | self.external_queue.reject('delivery_tag') 120 | 121 | def test_does_not_crash_on_context_management(self): 122 | queue = self.external_queue 123 | with queue as q: 124 | self.assertEqual(q, queue) 125 | # Do nothing to force close without connection 126 | 127 | def test_subscribe_declares_queue(self): 128 | with patch.object(self.external_queue, 'channel') as mock_channel: 129 | self.external_queue._subscribe() 130 | 131 | self.assertEqual(mock_channel.queue_declare.call_args_list, [call(queue=self.external_queue.default_queue_name, durable=False, exclusive=True, auto_delete=True)]) 132 | 133 | def test_subscribe_binds_to_queue(self): 134 | with patch.object(self.external_queue, 'channel') as mock_channel: 135 | self.external_queue._subscribe() 136 | 137 | self.assertEqual(mock_channel.queue_bind.call_args_list, [call(exchange='notifications', queue=self.external_queue.default_queue_name, routing_key=self.external_queue.subscription)]) 138 | 139 | def test_subscribe_sets_subscription(self): 140 | with patch.object(self.external_queue, '_connect'): 141 | self.external_queue.subscribe('routing_key') 142 | 143 | self.assertEqual('routing_key', self.external_queue.subscription) 144 | 145 | def test_subscribe_connects_to_server(self): 146 | with patch.object(self.external_queue, '_connect') as mock_connect: 147 | self.external_queue.subscribe('routing_key') 148 | 149 | self.assertEqual(mock_connect.call_count, 1) 150 | 151 | def test_message_callback_parses_message(self): 152 | def callback(message=None, delivery_tag=None): 153 | return message 154 | 155 | def stub_basic_consume(queue_name, callback): 156 | callback(self.message) 157 | 158 | mock_channel = Mock() 159 | mock_channel.basic_consume = stub_basic_consume 160 | 161 | with patch.object(self.external_queue, '_try'), \ 162 | patch.object(self.external_queue, '_parse_message', \ 163 | return_value=self.message) as mock_parse_message: 164 | self.external_queue.channel = mock_channel 165 | self.external_queue.setup_consumer(callback, ['queue_1']) 166 | 167 | self.assertEqual(mock_parse_message.call_args_list, [call(self.message)]) 168 | 169 | 170 | class ExceptionsTest(unittest.TestCase): 171 | def test_serializationerror_message(self): 172 | excp = Exception('Some random error') 173 | 174 | serialization_error = SerializationError(excp) 175 | 176 | self.assertEqual('{}'.format(serialization_error), 'Some random error') 177 | --------------------------------------------------------------------------------