Have manage.py queue_depths take multiple queues and return their depths

Author: jordaneb

django_dbq/management/commands/queue_depth.py

@@ -7,14 +7,17 @@ class Command(BaseCommand):
     help = "Print the current depth of the given queue"
 
     def add_arguments(self, parser):
-        parser.add_argument("queue_name", nargs="?", default="default", type=str)
+        parser.add_argument("queue_name", nargs="*", default=["default"], type=str)
 
     def handle(self, *args, **options):
-        queue_name = options["queue_name"]
+        queue_names = options["queue_name"]
         queue_depths = Job.get_queue_depths()
 
         self.stdout.write(
-            "queue_name={queue_name} queue_depth={depth}".format(
-                queue_name=queue_name, depth=queue_depths.get(queue_name, 0),
+            " ".join(
+                [
+                    f"{queue_name}={queue_depths.get(queue_name, 0)}"
+                    for queue_name in queue_names
+                ]
             )
         )

django_dbq/tests.py

@@ -117,13 +117,32 @@ def test_queue_depth(self):
         stdout = StringIO()
         call_command("queue_depth", stdout=stdout)
         output = stdout.getvalue()
-        self.assertEqual(output.strip(), "queue_name=default queue_depth=2")
+        self.assertEqual(output.strip(), "default=2")
+
+    def test_queue_depth_multiple_queues(self):
+
+        Job.objects.create(name="testjob", state=Job.STATES.FAILED)
+        Job.objects.create(name="testjob", state=Job.STATES.NEW)
+        Job.objects.create(name="testjob", state=Job.STATES.FAILED)
+        Job.objects.create(name="testjob", state=Job.STATES.COMPLETE)
+        Job.objects.create(name="testjob", state=Job.STATES.READY)
+        Job.objects.create(
+            name="testjob", queue_name="testqueue", state=Job.STATES.READY
+        )
+        Job.objects.create(
+            name="testjob", queue_name="testqueue", state=Job.STATES.READY
+        )
+
+        stdout = StringIO()
+        call_command("queue_depth", queue_name=("default", "testqueue",), stdout=stdout)
+        output = stdout.getvalue()
+        self.assertEqual(output.strip(), "default=2 testqueue=2")
 
     def test_queue_depth_for_queue_with_zero_jobs(self):
         stdout = StringIO()
-        call_command("queue_depth", queue_name="otherqueue", stdout=stdout)
+        call_command("queue_depth", queue_name=("otherqueue",), stdout=stdout)
         output = stdout.getvalue()
-        self.assertEqual(output.strip(), "queue_name=otherqueue queue_depth=0")
+        self.assertEqual(output.strip(), "otherqueue=0")
 
 
 @freezegun.freeze_time()