Add command line option for scheduler burst mode

rq_retry_scheduler/cli/scheduler.py

@@ -30,11 +30,15 @@ def get_arguments(args=None):
         default=os.environ.get('RQ_REDIS_PASSWORD'))
 
     parser.add_argument(
-        '--url', '-u',
+        '-u', '--url',
         default=os.environ.get('RQ_REDIS_URL'),
         help='URL describing Redis connection details. '
              'Overrides other connection arguments if supplied.')
 
+    parser.add_argument(
+        '-b', '--burst', action='store_true',
+        help='Burst mode. Move any jobs and quit')
+
     parser.add_argument(
         '-i', '--interval', help='Scheduler polling interval (in seconds)',
         default=10.0, type=float)
@@ -67,4 +71,4 @@ def main():
     setup_logging(args)
     connection = get_redis(args)
     scheduler = Scheduler(connection=connection, interval=args.interval)
-    scheduler.run()
+    scheduler.run(args.burst)

rq_retry_scheduler/scheduler.py

@@ -68,12 +68,16 @@ def enqueue_jobs(self):
             queue.enqueue_job(job)
             self.remove_job(job.id)
 
-    def run(self):
+    def run(self, burst=False):
         self.log.info('Starting RQ Retry Scheduler..')
 
         try:
             while True:
                 self.enqueue_jobs()
+
+                if burst:
+                    break
+
                 time.sleep(self.interval)
         finally:
             pass

tests/test_cli_main.py

@@ -29,15 +29,15 @@ def test_setup_logging():
 
 
 def test_main(mock):
-    args = Namespace(url='redis://localhost/15', interval=5)
+    args = Namespace(url='redis://localhost/15', interval=5, burst=False)
 
     mock.patch.object(scheduler, 'get_arguments', return_value=args)
     init = mock.spy(Scheduler, '__init__')
     run = mock.patch.object(Scheduler, 'run')
 
     scheduler.main()
 
-    assert run.called
+    run.assert_called_with(False)
 
     assert init.call_args[1]['interval'] == args.interval
 
@@ -49,3 +49,10 @@ def test_get_arguments():
     assert args.port == 6379
     assert args.url is None
     assert args.interval == 10.0
+    assert args.burst is False
+
+
+def test_burst_flag():
+    fake_arguments = ['-b']
+    args = scheduler.get_arguments(fake_arguments)
+    assert args.burst is True

tests/test_scheduler.py

@@ -139,3 +139,11 @@ def test_run(scheduler, mock):
         scheduler.run()
 
     assert enqueue_jobs.call_count == len(side_effects)
+
+
+def test_run_burst(scheduler, mock):
+    sleep = mock.patch.object(time, 'sleep')
+
+    scheduler.run(True)
+
+    assert not sleep.called