job decorator allow on_failure and on_success parameters (#1626) (#1628)

Co-authored-by: neptali <neptali@plerk.io>

rq/decorators.py

@@ -17,7 +17,8 @@ class job:  # noqa
     def __init__(self, queue, connection=None, timeout=None,
                  result_ttl=DEFAULT_RESULT_TTL, ttl=None,
                  queue_class=None, depends_on=None, at_front=None, meta=None,
-                 description=None, failure_ttl=None, retry=None):
+                 description=None, failure_ttl=None, retry=None, on_failure=None, 
+                 on_success=None):
         """A decorator that adds a ``delay`` method to the decorated function,
         which in turn creates a RQ job when called. Accepts a required
         ``queue`` argument that can be either a ``Queue`` instance or a string
@@ -41,6 +42,8 @@ class job:  # noqa
         self.description = description
         self.failure_ttl = failure_ttl
         self.retry = retry
+        self.on_success = on_success
+        self.on_failure = on_failure
 
     def __call__(self, f):
         @wraps(f)
@@ -65,6 +68,6 @@ class job:  # noqa
                                       timeout=self.timeout, result_ttl=self.result_ttl,
                                       ttl=self.ttl, depends_on=depends_on, job_id=job_id, at_front=at_front,
                                       meta=self.meta, description=self.description, failure_ttl=self.failure_ttl,
-                                      retry=self.retry)
+                                      retry=self.retry, on_failure=self.on_failure, on_success=self.on_success)
         f.delay = delay
         return f

tests/test_decorator.py

@@ -152,6 +152,43 @@ class TestDecorator(RQTestCase):
         self.assertEqual(baz_job.dependency, bar_job)
         self.assertEqual(baz_job.dependency.id, bar_job.id)
 
+    def test_decorator_accepts_on_failure_function_as_argument(self):
+        """Ensure that passing in on_failure function to the decorator sets the
+        correct on_failure function on the job.
+        """ 
+        # Only functions and builtins are supported as callback
+        @job('default', on_failure=Job.fetch)
+        def foo():
+            return 'Foo'
+        with self.assertRaises(ValueError):
+            result = foo.delay()
+
+        @job('default', on_failure=print)
+        def hello():
+            return 'Hello'
+        result = hello.delay()
+        result_job = Job.fetch(id=result.id, connection=self.testconn)
+        self.assertEqual(result_job.failure_callback, print)
+
+
+    def test_decorator_accepts_on_success_function_as_argument(self):
+        """Ensure that passing in on_failure function to the decorator sets the
+        correct on_success function on the job.
+        """
+        # Only functions and builtins are supported as callback
+        @job('default', on_failure=Job.fetch)
+        def foo():
+            return 'Foo'
+        with self.assertRaises(ValueError):
+            result = foo.delay()
+            
+        @job('default', on_success=print)
+        def hello():
+            return 'Hello'
+        result = hello.delay()
+        result_job = Job.fetch(id=result.id, connection=self.testconn)
+        self.assertEqual(result_job.success_callback, print)
+
     @mock.patch('rq.queue.resolve_connection')
     def test_decorator_connection_laziness(self, resolve_connection):
         """Ensure that job decorator resolve connection in `lazy` way """