From 141278bb42b36389b94c2d99b310a538004d168b Mon Sep 17 00:00:00 2001
From: Selwin Ong <selwin.ong@gmail.com>
Date: Fri, 30 May 2014 08:49:55 +0700
Subject: [PATCH] Queue.__init__ supports job_class argument.

---
 rq/queue.py         | 9 +++++++--
 tests/test_queue.py | 9 +++++++++
 2 files changed, 16 insertions(+), 2 deletions(-)

diff --git a/rq/queue.py b/rq/queue.py
index 4296001..baac6ca 100644
--- a/rq/queue.py
+++ b/rq/queue.py
@@ -6,7 +6,7 @@ import uuid
 
 from .connections import resolve_connection
 from .job import Job, Status
-from .utils import utcnow
+from .utils import import_attribute, utcnow
 
 from .exceptions import (DequeueTimeout, InvalidJobOperationError,
                          NoSuchJobError, UnpickleError)
@@ -55,7 +55,7 @@ class Queue(object):
         return cls(name, connection=connection)
 
     def __init__(self, name='default', default_timeout=None, connection=None,
-                 async=True):
+                 async=True, job_class=None):
         self.connection = resolve_connection(connection)
         prefix = self.redis_queue_namespace_prefix
         self.name = name
@@ -63,6 +63,11 @@ class Queue(object):
         self._default_timeout = default_timeout
         self._async = async
 
+        if job_class is not None:
+            if isinstance(job_class, string_types):
+                job_class = import_attribute(job_class)
+            self.job_class = job_class
+
     @property
     def key(self):
         """Returns the Redis key for this Queue."""
diff --git a/tests/test_queue.py b/tests/test_queue.py
index 1d467a6..d368590 100644
--- a/tests/test_queue.py
+++ b/tests/test_queue.py
@@ -12,6 +12,10 @@ from tests.fixtures import (div_by_zero, echo, Number, say_hello,
                             some_calculation)
 
 
+class CustomJob(Job):
+    pass
+
+
 class TestQueue(RQTestCase):
     def test_create_queue(self):
         """Creating queues."""
@@ -437,3 +441,8 @@ class TestFailedQueue(RQTestCase):
         q = Queue(async=False)
         job = q.enqueue(some_calculation, args=(2, 3))
         self.assertEqual(job.return_value, 6)
+
+    def test_custom_job_class(self):
+        """Ensure custom job class assignment works as expected."""
+        q = Queue(job_class=CustomJob)
+        self.assertEqual(q.job_class, CustomJob)