def test_buffer_publishes_with_confirm_mode_on(self): """ Verify buffering publishes, with confirm mode on, and then starting the producer results in a correct number of messages being sent successfully, measured by counting and verifying the Basic.Acks received from the broker. """ producer = RMQProducer() event = threading.Event() key_dict = dict() def on_confirm(confirm): if not isinstance(confirm, ConfirmModeOK): key_dict.pop(confirm) if len(key_dict.keys()) == 0: event.set() # Activate confirm mode producer.activate_confirm_mode(on_confirm) # Buffer a few publishes key_dict[producer.publish( b"body", exchange_params=ExchangeParams("exchange_direct"))] = False key_dict[producer.publish( b"body", exchange_params=ExchangeParams("exchange_direct"))] = False key_dict[producer.publish( b"body", exchange_params=ExchangeParams("exchange_direct"))] = False self.assertEqual(len(producer._unacked_publishes.keys()), 0) self.assertEqual(len(producer._buffered_messages), 3) # Start the producer producer.start() self.assertTrue(started(producer)) # Await all confirms event.wait(timeout=1.0) self.assertEqual(len(producer._unacked_publishes.keys()), 0) # Stop the extra producer producer.stop()
class IntegrationTest(unittest.TestCase): def setUp(self) -> None: self.consumer = RMQConsumer() self.consumer.start() self.assertTrue(started(self.consumer)) self.producer = RMQProducer() self.producer.start() self.assertTrue(started(self.producer)) def tearDown(self) -> None: self.consumer.stop() self.producer.stop() def test_stop_consumer(self): """ Verify RMQConsumer, when stopped, shuts down completely and releases allocated resources. """ # Run test self.consumer.stop() thread_count = len(threading.enumerate()) # Only MainThread should still be running. self.assertEqual(2, thread_count) def test_consume_from_queue(self): """Verify RMQConsumer can consume from a queue.""" msg_received = None event = threading.Event() def on_msg(msg): event.set() nonlocal msg_received msg_received = msg # Consume consume_key = self.consumer.consume(ConsumeParams(on_msg), queue_params=QueueParams("queue")) self.assertEqual(consume_key, "queue") # Wait for ConsumeOK event.wait(timeout=1.0) self.assertTrue(isinstance(msg_received, ConsumeOK)) # Send a message and verify it is delivered OK event.clear() # clear to re-use event self.producer.publish(b"body", queue_params=QueueParams("queue")) event.wait(timeout=1.0) self.assertEqual(msg_received, b"body") def test_consume_from_direct_exchange(self): """ Verify RMQConsumer can consume from a direct exchange and routing key. """ msg_received = None event = threading.Event() def on_msg(msg): event.set() nonlocal msg_received msg_received = msg # Consume consume_key = self.consumer.consume( ConsumeParams(on_msg), exchange_params=ExchangeParams("exchange"), routing_key="bla.bla") self.assertEqual(consume_key, "exchange|bla.bla") # Wait for ConsumeOK event.wait(timeout=1.0) self.assertTrue(isinstance(msg_received, ConsumeOK)) # Send a message and verify it is delivered OK event.clear() self.producer.publish(b"body", exchange_params=ExchangeParams("exchange"), routing_key="bla.bla") event.wait(timeout=1.0) self.assertEqual(msg_received, b"body") def test_consume_from_fanout_exchange(self): """ Verify RMQConsumer can consume from a fanout exchange. """ msg_received = None event = threading.Event() def on_msg(msg): event.set() nonlocal msg_received msg_received = msg # Consume consume_key = self.consumer.consume( ConsumeParams(on_msg), exchange_params=ExchangeParams("exchange_fanout", exchange_type=ExchangeType.fanout)) self.assertEqual(consume_key, "exchange_fanout") # Wait for ConsumeOK event.wait(timeout=1.0) self.assertTrue(isinstance(msg_received, ConsumeOK)) # Send a message and verify it is delivered OK event.clear() self.producer.publish(b"body", exchange_params=ExchangeParams( "exchange_fanout", exchange_type=ExchangeType.fanout)) event.wait(timeout=1.0) self.assertEqual(msg_received, b"body") def test_consume_from_exchange_and_queue(self): """ Verify RMQConsumer can consume from an exchange, binding it to a specific queue. """ msg_received = None event = threading.Event() def on_msg(msg): event.set() nonlocal msg_received msg_received = msg # Consume consume_key = self.consumer.consume( ConsumeParams(on_msg), exchange_params=ExchangeParams("exchange_fanout", exchange_type=ExchangeType.fanout), queue_params=QueueParams("queue_fanout_receiver")) self.assertEqual(consume_key, "queue_fanout_receiver|exchange_fanout") # Wait for ConsumeOK event.wait(timeout=1.0) self.assertTrue(isinstance(msg_received, ConsumeOK)) # Send a message and verify it is delivered OK event.clear() self.producer.publish(b"body", exchange_params=ExchangeParams( "exchange_fanout", exchange_type=ExchangeType.fanout)) event.wait(timeout=1.0) self.assertEqual(msg_received, b"body") def test_consume_from_exchange_routing_key_and_queue(self): """ Verify RMQConsumer can consume from an exchange, binding it to a specific queue. """ msg_received = None event = threading.Event() def on_msg(msg): event.set() nonlocal msg_received msg_received = msg # Consume consume_key = self.consumer.consume( ConsumeParams(on_msg), exchange_params=ExchangeParams("exchange_direct"), routing_key="fish", queue_params=QueueParams("queue_direct_receiver")) self.assertEqual(consume_key, "queue_direct_receiver|exchange_direct|fish") # Wait for ConsumeOK event.wait(timeout=1.0) self.assertTrue(isinstance(msg_received, ConsumeOK)) # Send a message and verify it is delivered OK event.clear() self.producer.publish( b"body", exchange_params=ExchangeParams("exchange_direct"), routing_key="fish") event.wait(timeout=1.0) self.assertEqual(msg_received, b"body") def test_confirm_mode(self): """ Verify confirm mode pushes publish_keys to users on successful delivery. """ msg_received = None event = threading.Event() def on_confirm(confirm): event.set() nonlocal msg_received msg_received = confirm # Activate confirm mode self.producer.activate_confirm_mode(on_confirm) # Wait for ConfirmModeOK event.wait(timeout=1.0) self.assertTrue(isinstance(msg_received, ConfirmModeOK)) # Send a message and verify it is delivered OK event.clear() publish_key = self.producer.publish( b"body", exchange_params=ExchangeParams("exchange_direct"), ) event.wait(timeout=1.0) self.assertEqual(msg_received, publish_key) self.assertEqual(len(self.producer._unacked_publishes.keys()), 0) def test_buffer_publishes_with_confirm_mode_on(self): """ Verify buffering publishes, with confirm mode on, and then starting the producer results in a correct number of messages being sent successfully, measured by counting and verifying the Basic.Acks received from the broker. """ producer = RMQProducer() event = threading.Event() key_dict = dict() def on_confirm(confirm): if not isinstance(confirm, ConfirmModeOK): key_dict.pop(confirm) if len(key_dict.keys()) == 0: event.set() # Activate confirm mode producer.activate_confirm_mode(on_confirm) # Buffer a few publishes key_dict[producer.publish( b"body", exchange_params=ExchangeParams("exchange_direct"))] = False key_dict[producer.publish( b"body", exchange_params=ExchangeParams("exchange_direct"))] = False key_dict[producer.publish( b"body", exchange_params=ExchangeParams("exchange_direct"))] = False self.assertEqual(len(producer._unacked_publishes.keys()), 0) self.assertEqual(len(producer._buffered_messages), 3) # Start the producer producer.start() self.assertTrue(started(producer)) # Await all confirms event.wait(timeout=1.0) self.assertEqual(len(producer._unacked_publishes.keys()), 0) # Stop the extra producer producer.stop() @classmethod def tearDownClass(cls) -> None: """ Need general teardown since not wanting to introduce waits part of test cases. Need to purge queues as consumer/producers are stopped before acking messages, sometimes. This teardown enabled re-testability without getting messages sent by a previous run. """ client = RMQConsumer() client.start() assert started(client) client._channel.queue_purge("queue_fanout_receiver") client._channel.queue_purge("queue") client.stop()
class TestConfirmMode(unittest.TestCase): """ Test the confirm mode functionality of RMQProducer. """ @patch("rabbitmq_client.producer.RMQProducer.start") def setUp(self, _connection_start) -> None: """Setup to run before each test case.""" self.producer = RMQProducer() self.producer.start() self.producer.on_ready() # Fake connection getting ready self.producer.confirm_delivery = Mock() self.producer.declare_exchange = Mock() self.producer.basic_publish = Mock() def test_activate_confirm_mode(self): """ Verify that confirm mode, when activated, yields the expected state. """ # Prep confirm_mode_started = False def notify_callback(confirm): nonlocal confirm_mode_started if isinstance(confirm, ConfirmModeOK): confirm_mode_started = True # Test self.producer.activate_confirm_mode(notify_callback) self.producer.confirm_delivery.assert_called_with( self.producer.on_delivery_confirmed, callback=self.producer.on_confirm_select_ok) self.assertFalse(self.producer._confirm_mode_active) self.producer.on_confirm_select_ok(Mock()) # Assert self.assertTrue(confirm_mode_started) self.assertTrue(self.producer._confirm_mode_active) def test_confirm_mode_publish_key_returned_by_publish(self): """ Verify that a publish key is generated and returned by publish if confirm mode is activated. """ # Prep publish_key = None def on_confirm(delivery_notification): if delivery_notification == publish_key: return elif isinstance(delivery_notification, ConfirmModeOK): return self.fail("delivery notification did not match the generated " "publish key") self.producer.activate_confirm_mode(on_confirm) self.producer.confirm_delivery.assert_called_with( self.producer.on_delivery_confirmed, callback=self.producer.on_confirm_select_ok) self.producer.on_confirm_select_ok(None) exchange_params = ExchangeParams("exchange") # Test publish_key = self.producer.publish(b"body", exchange_params=exchange_params) self.producer.declare_exchange.assert_called() self.producer.on_exchange_declared(b"body", exchange_params, None, publish_key=publish_key) self.producer.basic_publish.assert_called_with( b"body", exchange=exchange_params.exchange, routing_key="", publish_params=ANY) # 1 is the delivery tag self.assertEqual(self.producer._unacked_publishes[1], publish_key) frame = Mock() ack = Basic.Ack() ack.delivery_tag = 1 frame.method = ack self.producer.on_delivery_confirmed(frame) self.assertEqual(len(self.producer._unacked_publishes.keys()), 0) def test_on_ready_initiates_confirm_mode_again(self): """ Verify that after connectivity issues, on_ready will re-start confirm- mode. """ # Prep self.producer.activate_confirm_mode(lambda _:...) # Test self.producer.on_ready() # Assert self.producer.confirm_delivery.assert_called_with( self.producer.on_delivery_confirmed, callback=self.producer.on_confirm_select_ok) def test_delivery_not_acked_leads_to_delivery_error(self): """ Verify that if a delivery confirmation is NOT an ack, a delivery error is propagated to the delivery notify callback. """ # Prep got_error = False def on_confirm(error): nonlocal got_error got_error = True if isinstance(error, DeliveryError): self.assertEqual("123", error.publish_key) self.producer.activate_confirm_mode(on_confirm) self.producer._unacked_publishes[1] = "123" # Test frame = Mock() frame.method = Basic.Nack() frame.method.delivery_tag = 1 self.producer.on_delivery_confirmed(frame) # Assert self.assertTrue(got_error) def test_activate_confirm_mode_idempotent(self): """ Verify that calling activate_confirm_mode more than once has no effect, and will yield the same result as calling it once. """ # Prep def on_confirm(): ... # Test + assert self.producer.activate_confirm_mode(on_confirm) self.producer.activate_confirm_mode(on_confirm) self.producer.confirm_delivery.assert_called_once() def test_confirm_mode_delays_buffer_until_confirm_mode_ok(self): """ Verify that, if confirm mode has been activated, buffered publishes are not sent on 'on_ready', but rather 'on_confirm_select_ok'. """ # Prep self.producer.activate_confirm_mode(lambda _:...) self.producer.on_close() exchange_params = ExchangeParams("exchange") # Test + assertions _pub_key = self.producer.publish(b"body", exchange_params=exchange_params) self.assertNotEqual(_pub_key, None) self.producer.declare_exchange.assert_not_called() self.assertEqual(len(self.producer._buffered_messages), 1) self.producer.on_ready() self.producer.declare_exchange.assert_not_called() self.assertEqual(len(self.producer._buffered_messages), 1) self.producer.on_confirm_select_ok(Mock()) self.producer.declare_exchange.assert_called() self.assertEqual(len(self.producer._buffered_messages), 0) def test_publish_before_confirm_mode_ok_buffers_the_publish(self): """ Verify that any call to publish before confirm mode ok (when confirm mode has been activated) leads to the call being buffered and delayed until 'on_confirm_select_ok'. """ # Prep self.producer.activate_confirm_mode(lambda _:...) exchange_params = ExchangeParams("exchange") # Test + assertions _pub_key = self.producer.publish(b"body", exchange_params=exchange_params) self.assertNotEqual(_pub_key, None) self.producer.declare_exchange.assert_not_called() self.assertEqual(len(self.producer._buffered_messages), 1) self.producer.on_confirm_select_ok(Mock()) self.producer.declare_exchange.assert_called() self.assertEqual(len(self.producer._buffered_messages), 0)
def stop(*args): producer.stop() sys.exit(0) signal.signal(signal.SIGINT, stop) producer.publish(b"queue publish 1", queue_params=QueueParams("queue")) producer.publish(b"queue publish 2", queue_params=QueueParams("queue")) producer.publish(b"queue publish 3", queue_params=QueueParams("queue")) producer.publish(b"exchange publish 1", exchange_params=ExchangeParams("direct"), routing_key="rkey") producer.publish(b"exchange publish 2", exchange_params=ExchangeParams( "fanout", exchange_type=ExchangeType.fanout)) time.sleep(0.3) producer.activate_confirm_mode(lambda x: print(x)) producer.publish(b"exchange publish 2", exchange_params=ExchangeParams( "fanout", exchange_type=ExchangeType.fanout)) producer.publish(b"exchange publish 2", exchange_params=ExchangeParams( "fanout", exchange_type=ExchangeType.fanout)) producer.publish(b"exchange publish 2", exchange_params=ExchangeParams( "fanout", exchange_type=ExchangeType.fanout)) threading.Event().wait()