Add support for deleting queues after flushing them

dramatiq/brokers/rabbitmq.py

@@ -382,21 +382,29 @@ def get_queue_message_counts(self, queue_name):
             xq_queue_response.method.message_count,
         )
 
-    def flush(self, queue_name):
+    def flush(self, queue_name: str, delete_queue: bool = False):
         """Drop all the messages from a queue.
 
         Parameters:
           queue_name(str): The queue to flush.
+          delete_queue(bool): If true, the queue will be deleted after it is purged.
         """
         for name in (queue_name, dq_name(queue_name), xq_name(queue_name)):
             if queue_name not in self.queues_pending:
+                self.logger.info("Purging queue %s", queue_name)
                 self.channel.queue_purge(name)
 
-    def flush_all(self):
+                if delete_queue:
+                    self.channel.queue_delete(name)
+
+    def flush_all(self, delete_queue: bool = False):
         """Drop all messages from all declared queues.
+
+        Parameters:
+          delete_queue(bool): If true, the queues will be deleted after they're flushed.
         """
         for queue_name in self.queues:
-            self.flush(queue_name)
+            self.flush(queue_name, delete_queue)
 
     def join(self, queue_name, min_successes=10, idle_time=100, *, timeout=None):
         """Wait for all the messages on the given queue to be

tests/conftest.py

@@ -70,7 +70,7 @@ def rabbitmq_broker():
     broker.emit_after("process_boot")
     dramatiq.set_broker(broker)
     yield broker
-    broker.flush_all()
+    broker.flush_all(delete_queue=True)
     broker.close()
 
 

tests/test_rabbitmq.py

@@ -503,3 +503,35 @@ def put():
 
     assert len(rabbitmq_broker.queues) == 1
     assert put.queue_name in rabbitmq_broker.queues
+
+
+def test_rabbitmq_flush_true_deletes_the_queue():
+    queue_name = f"flush_all_test_queue_{current_millis()}"
+    url = "amqp://%s:%[email protected]:5672" % (RABBITMQ_USERNAME, RABBITMQ_PASSWORD)
+    broker = RabbitmqBroker(url=url)
+    broker.declare_queue(queue_name, ensure=True)
+
+    broker.flush(queue_name, delete_queue=False)
+    assert broker.channel.queue_declare(queue=queue_name, passive=True)
+
+    broker.flush(queue_name, delete_queue=True)
+    with pytest.raises(pika.exceptions.ChannelClosedByBroker, match=r"NOT_FOUND - no queue"):
+        broker.channel.queue_declare(queue=queue_name, passive=True)
+
+
+def test_rabbitmq_flush_all_true_deletes_the_queue():
+    queues = set(f"flush_all_test_queue_{current_millis()}_{i}" for i in range(10))
+    url = "amqp://%s:%[email protected]:5672" % (RABBITMQ_USERNAME, RABBITMQ_PASSWORD)
+    broker = RabbitmqBroker(url=url)
+
+    for queue_name in queues:
+        broker.declare_queue(queue_name, ensure=True)
+
+    broker.flush_all(delete_queue=False)
+    for queue_name in queues:
+        assert broker.channel.queue_declare(queue=queue_name, passive=True)
+
+    broker.flush_all(delete_queue=True)
+    for queue_name in queues:
+        with pytest.raises(pika.exceptions.ChannelClosedByBroker, match=r"NOT_FOUND - no queue"):
+            broker.connection.channel().queue_declare(queue=queue_name, passive=True)