get_scheduler: allow passing a custom connection object (#661)

rq · May 1, 2024 · d33163d · d33163d
1 parent fa23214
commit d33163d
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 7 deletions.
diff --git a/django_rq/queues.py b/django_rq/queues.py
@@ -342,6 +342,7 @@ def get_scheduler(
         name: str = 'default',
         queue: Optional[DjangoRQ] = None,
         interval: int = 60,
+        connection: Optional[redis.Redis] = None,
     ) -> DjangoScheduler:
         """
         Returns an RQ Scheduler instance using parameters defined in
@@ -353,11 +354,14 @@ def get_scheduler(
         if isinstance(scheduler_class, str):
             scheduler_class = import_attribute(scheduler_class)
 
+        if connection is None:
+            connection = get_connection(name)
+
         if queue is None:
-            queue = get_queue(name)
+            queue = get_queue(name, connection=connection)
 
         return scheduler_class(
-            queue_name=name, interval=interval, queue=queue, job_class=queue.job_class, connection=get_connection(name)
+            queue_name=name, interval=interval, queue=queue, job_class=queue.job_class, connection=connection
         )
 
 except ImportError:

diff --git a/django_rq/tests/tests.py b/django_rq/tests/tests.py
@@ -681,8 +681,8 @@ def test_error(self):
         self.assertEqual(queue.count, 0)
 
 
+@skipIf(RQ_SCHEDULER_INSTALLED is False, 'RQ Scheduler not installed')
 class SchedulerTest(TestCase):
-    @skipIf(RQ_SCHEDULER_INSTALLED is False, 'RQ Scheduler not installed')
     def test_get_scheduler(self):
         """
         Ensure get_scheduler creates a scheduler instance with the right
@@ -696,7 +696,16 @@ def test_get_scheduler(self):
         self.assertEqual(connection_kwargs['port'], config['PORT'])
         self.assertEqual(connection_kwargs['db'], config['DB'])
 
-    @skipIf(RQ_SCHEDULER_INSTALLED is False, 'RQ Scheduler not installed')
+    def test_get_scheduler_custom_connection(self):
+        """
+        Ensure get_scheduler respects the `connection` argument.
+        """
+
+        with get_connection('test') as connection:
+            scheduler = get_scheduler('test', connection=connection)
+
+            self.assertIs(scheduler.connection, connection)
+
     @patch('django_rq.management.commands.rqscheduler.get_scheduler')
     @patch('django_rq.management.commands.rqscheduler.setup_loghandlers')
     def test_commandline_verbosity_affects_logging_level(self, setup_loghandlers_mock, get_scheduler_mock):
@@ -713,7 +722,6 @@ def test_commandline_verbosity_affects_logging_level(self, setup_loghandlers_moc
             setup_loghandlers_mock.assert_called_once_with(expected_level[verbosity])
 
     @override_settings(RQ={'SCHEDULER_CLASS': 'django_rq.tests.fixtures.DummyScheduler'})
-    @skipIf(RQ_SCHEDULER_INSTALLED is False, 'RQ Scheduler not installed')
     def test_scheduler_default_timeout(self):
         """
         Scheduler class customization.
@@ -722,7 +730,6 @@ def test_scheduler_default_timeout(self):
         self.assertIsInstance(scheduler, DummyScheduler)
 
     @override_settings(RQ={'AUTOCOMMIT': True})
-    @skipIf(RQ_SCHEDULER_INSTALLED is False, 'RQ Scheduler not installed')
     def test_scheduler_default_timeout(self):
         """
         Ensure scheduler respects DEFAULT_RESULT_TTL value for `result_ttl` param.
@@ -734,7 +741,6 @@ def test_scheduler_default_timeout(self):
         job.delete()
 
     @override_settings(RQ={'AUTOCOMMIT': True, 'DEFAULT_RESULT_TTL': 5432})
-    @skipIf(RQ_SCHEDULER_INSTALLED is False, 'RQ Scheduler not installed')
     def test_scheduler_default_result_ttl(self):
         """
         Ensure scheduler respects DEFAULT_RESULT_TTL value for `result_ttl` param.