Merge pull request #2 from dabapps/fix-worker

Fix worker

Author: j4mie

django_dbq/management/commands/worker.py

@@ -1,7 +1,8 @@
 from django.db import transaction
-from django.core.management.base import NoArgsCommand
+from django.core.management.base import BaseCommand, CommandError
 from django.utils.module_loading import import_by_path
 from django_dbq.models import Job
+from optparse import make_option
 from simplesignals.process import WorkerProcessBase
 from time import sleep
 import logging
@@ -10,6 +11,9 @@
 logger = logging.getLogger(__name__)
 
 
+DEFAULT_QUEUE_NAME = 'default'
+
+
 def process_job(queue_name):
     """This function grabs the next available job for a given queue, and runs its next task."""
 
@@ -64,17 +68,32 @@ def do_work(self):
         process_job(self.queue_name)
 
 
-class Command(NoArgsCommand):
+class Command(BaseCommand):
 
     help = "Run a queue worker process"
 
-    def handle_noargs(self, **options):
-        self.handle('default')
+    option_list = BaseCommand.option_list + (
+        make_option('--dry-run',
+            action='store_true',
+            dest='dry_run',
+            default=False,
+            help="Don't actually start the worker. Used for testing."),
+        )
 
     def handle(self, *args, **options):
+        if not args:
+            args = (DEFAULT_QUEUE_NAME,)
+
         if len(args) != 1:
             raise CommandError("Please supply a single queue job name")
 
-        logger.info("Starting job worker")
-        worker = Worker()
+        queue_name = args[0]
+
+        self.stdout.write("Starting job worker for queue \"%s\"" % queue_name)
+
+        worker = Worker(queue_name)
+
+        if options['dry_run']:
+            return
+
         worker.run()

django_dbq/tests.py

@@ -58,6 +58,23 @@ def test_errors_raised_correctly(self):
             call_command('create_job', 'some_other_job', stdout=StringIO())
 
 
+@override_settings(JOBS={'testjob': {'tasks': ['a']}})
+class WorkerManagementCommandTestCase(TestCase):
+
+    def test_worker_no_args(self):
+        stdout = StringIO()
+        call_command('worker', dry_run=True, stdout=stdout)
+        output = stdout.getvalue()
+        self.assertTrue('Starting job worker' in output)
+        self.assertTrue('default' in output)
+
+    def test_worker_with_queue_name(self):
+        stdout = StringIO()
+        call_command('worker', 'test_queue', dry_run=True, stdout=stdout)
+        output = stdout.getvalue()
+        self.assertTrue('test_queue' in output)
+
+
 @override_settings(JOBS={'testjob': {'tasks': ['a']}})
 class JobTestCase(TestCase):