class TestProducerBootstrap(TestBase):
def setup(self):
super().setup()
self.producer = Producer(self.channel, self.exchange_name)
self.queues = ['queue-1', 'queue-2', 'queue-3', 'queue-4']
queue_handler = patch('broker_rabbit.producer.QueueHandler')
self.queue_handler = queue_handler.start()
exchange_handler = patch('broker_rabbit.producer.ExchangeHandler')
self.exchange_handler = exchange_handler.start()
def test_should_setup_exchange(self):
# When
self.producer.bootstrap(self.queues)
# Then
channel = self.channel.get_channel()
self.exchange_handler.assert_called_once_with(channel,
self.exchange_name)
self.exchange_handler().setup_exchange.assert_called_once()
def test_should_setup_queue(self):
# When
self.producer.bootstrap(self.queues)
# Then
channel = self.channel.get_channel()
self.queue_handler.assert_called_once_with(channel, self.exchange_name)
assert 4 == self.queue_handler().setup_queue.call_count
def test_should_close_channel_at_the_end(self):
# When
self.producer.bootstrap(self.queues)
# Then
self.channel.close.assert_called_once()
def test_should_close_channel_at_the_end_while_error_occurred(self):
# Given
channel = self.channel.get_channel()
error_msg = 'connection not opened'
channel.open.side_effect = ConnectionNotOpenedError(error_msg)
# When
self.producer.bootstrap(self.queues)
# Then
self.channel.close.assert_called_once()
class TestProducerBootstrap(TestBase):
def setup(self):
super().setup()
self.producer = Producer(self.channel, self.exchange_name)
self.queues = ['queue-1', 'queue-2', 'queue-3', 'queue-4']
def test_should_setup_exchange(self, exchange_mock, queue_mock):
# When
self.producer.bootstrap(self.queues)
# Then
channel = self.channel.get_channel()
exchange_mock.assert_called_once_with(channel, self.exchange_name)
exchange_mock().setup_exchange.assert_called_once()
def test_should_setup_queue(self, exchange_mock, queue_mock):
# When
self.producer.bootstrap(self.queues)
# Then
channel = self.channel.get_channel()
queue_mock.assert_called_once_with(channel, self.exchange_name)
assert 4 == queue_mock().setup_queue.call_count
@pytest.mark.skip
def test_should_close_channel_at_the_end(self, exchange_mock, queue_mock):
# When
self.producer.bootstrap(self.queues)
# Then
channel = self.channel.get_channel()
channel.close.assert_called_once()
@pytest.mark.skip
def test_should_close_channel_at_the_end_while_error_occurred(
self, exchange_mock, queue_mock):
# Given
channel = self.channel.get_channel()
channel.open.side_effect = ConnectionNotOpenedYet(
'connection not opened')
# When
self.producer.bootstrap(self.queues)
# Then
channel.close.assert_called_once()
class BrokerRabbitMQ:
"""Message Broker based on RabbitMQ middleware"""
def __init__(self, app=None, queues=None):
"""
Create a new instance of Broker Rabbit by using
the given parameters to connect to RabbitMQ.
"""
self.app = app
self.queues = queues
self.connection_handler = None
self.producer = None
self.url = None
self.exchange_name = None
self.application_id = None
self.delivery_mode = None
self.on_message_callback = None
if self.app is not None:
self.init_app(self.app, self.queues)
def init_app(self, app, queues, on_message_callback):
""" Init the Broker by using the given configuration instead
default settings.
:param app: Current application context
:param list queues: Queues which the message will be post
:param callback on_message_callback: callback to execute when new
message is pulled to RabbitMQ
"""
if not hasattr(app, 'extensions'):
app.extensions = {}
if 'broker_rabbit' not in app.extensions:
app.extensions['broker_rabbit'] = self
else:
# Raise an exception if extension already initialized as
# potentially new configuration would not be loaded.
raise RuntimeError('Extension already initialized')
self.url = app.config.get('RABBIT_MQ_URL', DEFAULT_URL)
self.exchange_name = app.config.get('EXCHANGE_NAME', DEFAULT_EXCHANGE_NAME)
self.application_id = app.config.get('APPLICATION_ID', DEFAULT_APPLICATION_ID)
self.delivery_mode = app.config.get('DELIVERY_MODE', DEFAULT_DELIVERY_MODE)
self.queues = queues
self.on_message_callback = on_message_callback
# Open Connection to RabbitMQ
self.connection_handler = ConnectionHandler(self.url)
connection = self.connection_handler.get_current_connection()
# Setup default producer for broker_rabbit
self.producer = Producer(
connection, self.exchange_name,
self.application_id, self.delivery_mode)
self.producer.bootstrap(self.queues)
def send(self, queue, context={}):
"""Post the message to the correct queue with the given context
:param str queue: queue which to post the message
:param dict context: content of the message to post to RabbitMQ server
"""
if queue not in self.queues:
message = 'Queue ‘{name}‘ is not registered'.format(name=queue)
raise UnknownQueueError(message)
message = {
'created_at': datetime.utcnow().isoformat(),
'queue': queue,
'context': context
}
return self.producer.publish(queue, message)
class BrokerRabbitMQ:
"""Message Broker based on RabbitMQ middleware"""
def __init__(self, app=None):
"""
Create a new instance of Broker Rabbit by using
the given parameters to connect to RabbitMQ.
"""
self.app = app
self.connection_handler = None
self.producer = None
self.url = None
self.exchange_name = None
self.application_id = None
self.delivery_mode = None
self.queues = None
self.on_message_callback = None
def init_app(self, app, queues, on_message_callback):
""" Init the Broker by using the given configuration instead
default settings.
:param app: Current application context
:param list queues: Queues which the message will be post
:param callback on_message_callback: callback to execute when new
message is pulled to RabbitMQ
"""
self.url = app.config.get('RABBIT_MQ_URL', DEFAULT_URL)
self.exchange_name = app.config.get('EXCHANGE_NAME', DEFAULT_EXCHANGE)
self.application_id = app.config.get('APPLICATION_ID', DEFAULT_APP)
self.delivery_mode = app.config.get('DELIVERY_MODE', DEFAULT_DELIVERY)
self.queues = queues
self.on_message_callback = on_message_callback
# Open Connection to RabbitMQ
self.connection_handler = ConnectionHandler(self.url)
connection = self.connection_handler.get_current_connection()
# Setup default producer for broker_rabbit
channel = ProducerChannel(connection, self.application_id,
self.delivery_mode)
self.producer = Producer(channel, self.exchange_name)
self.producer.bootstrap(self.queues)
def send(self, queue, context={}):
"""Post the message to the correct queue with the given context
:param str queue: queue which to post the message
:param dict context: content of the message to post to RabbitMQ server
"""
if queue not in self.queues:
error_msg = f'Queue ‘{queue}‘ is not registered'
raise UnknownQueueError(error_msg)
message = {
'created_at': datetime.utcnow().isoformat(),
'queue': queue,
'context': context
}
return self.producer.publish(queue, message)
def list_queues(self):
"""List all available queue in the app"""
for queue in self.queues:
print('Queue name : `%s`' % queue)
def start(self, queue):
"""Start worker on a given queue
:param queue: the queue which you consume message for
"""
if queue not in self.queues:
raise RuntimeError(f'Queue with name`{queue}` not found')
worker = Worker(connection_handler=self.connection_handler,
message_callback=self.on_message_callback,
queue=queue)
print(f'Start consuming message on the queue `{queue}`')
worker.consume_message()