Add flush_queues

servicelayer/taskqueue.py

@@ -822,3 +822,18 @@ def queue_task(
                 channel.close()
         except pika.exceptions.ChannelWrongStateError:
             log.exception("Failed to explicitly close RabbitMQ channel.")
+
+
+def flush_queues(rmq_conn, redis_conn, queues):
+    try:
+        channel = rmq_conn.channel()
+        for queue in queues:
+            try:
+                channel.queue_purge(queue)
+            except ValueError:
+                logging.exception(f"Error while flushing the {queue} queue")
+        channel.close()
+    except pika.exceptions.AMQPError:
+        logging.exception("Error while flushing task queue")
+    for key in redis_conn.scan_iter(PREFIX + "*"):
+        redis_conn.delete(key)

tests/test_taskqueue.py

@@ -15,6 +15,7 @@
     get_rabbitmq_connection,
     dataset_from_collection_id,
     declare_rabbitmq_queue,
+    flush_queues,
 )
 from servicelayer.util import unpack_datetime
 
@@ -190,9 +191,9 @@ def did_nack():
 
 
 def test_get_priority_bucket():
-    # flush_queue()
     redis = get_fakeredis()
     rmq = get_rabbitmq_connection()
+    flush_queues(rmq, redis, ["index"])
     collection = Mock(id=1)
 
     assert get_task_count(collection, redis) == 0