diff --git a/rq/decorators.py b/rq/decorators.py index 309c09c..1f3a78e 100644 --- a/rq/decorators.py +++ b/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 diff --git a/tests/test_decorator.py b/tests/test_decorator.py index 5e846db..32eaa91 100644 --- a/tests/test_decorator.py +++ b/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 """