add run_check (#116)

* add run_check, fix #115

* fix job results

* better exception message

Author: samuelcolvin

HISTORY.rst

@@ -3,6 +3,10 @@
 History
 -------
 
+v0.16.0a4 (unreleased)
+......................
+* add ``Worker.run_check``, fix #115
+
 v0.16.0a3 (2010-03-12)
 ......................
 * fix ``Worker`` with custom redis settings

arq/jobs.py

@@ -47,10 +47,10 @@ async def result(self, timeout: Optional[float] = None, *, pole_delay: float = 0
             info = await self.result_info()
             if info:
                 result = info['result']
-                if isinstance(result, Exception):
-                    raise result
-                else:
+                if info['success']:
                     return result
+                else:
+                    raise result
             if timeout is not None and delay > timeout:
                 raise asyncio.TimeoutError()
 
@@ -81,14 +81,15 @@ async def result_info(self) -> Optional[Dict[str, Any]]:
         """
         v = await self._redis.get(result_key_prefix + self.job_id, encoding=None)
         if v:
-            enqueue_time_ms, job_try, function, args, kwargs, result, start_time_ms, finish_time_ms = pickle.loads(v)
+            enqueue_time_ms, job_try, function, args, kwargs, s, r, start_time_ms, finish_time_ms = pickle.loads(v)
             return dict(
                 enqueue_time=ms_to_datetime(enqueue_time_ms),
                 job_try=job_try,
                 function=function,
                 args=args,
                 kwargs=kwargs,
-                result=result,
+                result=r,
+                success=s,
                 start_time=ms_to_datetime(start_time_ms),
                 finish_time=ms_to_datetime(finish_time_ms),
             )

arq/version.py

@@ -2,4 +2,4 @@
 
 __all__ = ['VERSION']
 
-VERSION = StrictVersion('0.16a3')
+VERSION = StrictVersion('0.16a4')

arq/worker.py

@@ -100,6 +100,22 @@ def __str__(self):
         return repr(self)
 
 
+class FailedJobs(RuntimeError):
+    def __init__(self, count, job_results):
+        self.count = count
+        self.job_results = job_results
+
+    def __str__(self):
+        if self.count == 1 and self.job_results:
+            exc = self.job_results[0]['result']
+            return f'1 job failed "{exc.__class__.__name__}: {exc}"'
+        else:
+            return f'{self.count} jobs failed'
+
+    def __repr__(self):
+        return f'<{str(self)}>'
+
+
 class Worker:
     """
     Main class for running jobs.
@@ -174,7 +190,10 @@ def __init__(
         self._add_signal_handler(signal.SIGTERM, self.handle_sig)
         self.on_stop = None
 
-    def run(self):
+    def run(self) -> None:
+        """
+        Sync function to run the worker, finally closes worker connections.
+        """
         self.main_task = self.loop.create_task(self.main())
         try:
             self.loop.run_until_complete(self.main_task)
@@ -184,10 +203,27 @@ def run(self):
         finally:
             self.loop.run_until_complete(self.close())
 
-    async def async_run(self):
+    async def async_run(self) -> None:
+        """
+        Asynchronously run the worker, does not close connections. Useful when testing.
+        """
         self.main_task = self.loop.create_task(self.main())
         await self.main_task
 
+    async def run_check(self) -> int:
+        """
+        Run :func:`arq.worker.Worker.async_run`, check for failed jobs and raise :class:`arq.worker.FailedJobs`
+        if any jobs have failed.
+
+        :return: number of completed jobs
+        """
+        await self.async_run()
+        if self.jobs_failed:
+            failed_job_results = [r for r in await self.pool.all_job_results() if not r['success']]
+            raise FailedJobs(self.jobs_failed, failed_job_results)
+        else:
+            return self.jobs_complete
+
     async def main(self):
         if self.pool is None:
             self.pool = await create_pool(self.redis_settings)
@@ -293,6 +329,7 @@ async def run_job(self, job_id, score):  # noqa: C901
         }
         ctx = {**self.ctx, **job_ctx}
         start_ms = timestamp_ms()
+        success = False
         try:
             s = args_to_string(args, kwargs)
             extra = f' try={job_try}' if job_try > 1 else ''
@@ -330,6 +367,7 @@ async def run_job(self, job_id, score):  # noqa: C901
                 finish = True
                 self.jobs_failed += 1
         else:
+            success = True
             finished_ms = timestamp_ms()
             logger.info('%6.2fs ← %s ● %s', (finished_ms - start_ms) / 1000, ref, result_str)
             finish = True
@@ -338,7 +376,7 @@ async def run_job(self, job_id, score):  # noqa: C901
         result_timeout_s = self.keep_result_s if function.keep_result_s is None else function.keep_result_s
         result_data = None
         if result is not no_result and result_timeout_s > 0:
-            d = enqueue_time_ms, job_try, function_name, args, kwargs, result, start_ms, finished_ms
+            d = enqueue_time_ms, job_try, function_name, args, kwargs, success, result, start_ms, finished_ms
             try:
                 result_data = pickle.dumps(d)
             except AttributeError:

tests/test_jobs.py

@@ -48,6 +48,7 @@ async def foobar(ctx, *args, **kwargs):
         'function': 'foobar',
         'args': (1, 2),
         'kwargs': {'c': 3},
+        'success': True,
         'result': 42,
         'start_time': CloseToNow(),
         'finish_time': CloseToNow(),
@@ -61,6 +62,7 @@ async def foobar(ctx, *args, **kwargs):
             'function': 'foobar',
             'args': (1, 2),
             'kwargs': {'c': 3},
+            'success': True,
             'result': 42,
             'start_time': CloseToNow(),
             'finish_time': CloseToNow(),

tests/test_worker.py

@@ -4,17 +4,22 @@
 import signal
 from unittest.mock import MagicMock
 
+import pytest
 from aioredis import create_redis_pool
 
 from arq.connections import ArqRedis
 from arq.constants import health_check_key, job_key_prefix
-from arq.worker import Retry, Worker, async_check_health, check_health, func, run_worker
+from arq.worker import FailedJobs, Retry, Worker, async_check_health, check_health, func, run_worker
 
 
 async def foobar(ctx):
     return 42
 
 
+async def fails(ctx):
+    raise TypeError('my type error')
+
+
 def test_no_jobs(arq_redis: ArqRedis, loop):
     class Settings:
         functions = [func(foobar, name='foobar')]
@@ -287,3 +292,49 @@ async def test_remain_keys_no_results(arq_redis: ArqRedis, worker):
     worker: Worker = worker(functions=[func(foobar, keep_result=0)])
     await worker.main()
     assert sorted(await arq_redis.keys('*')) == ['arq:health-check']
+
+
+async def test_run_check_passes(arq_redis: ArqRedis, worker):
+    await arq_redis.enqueue_job('foobar')
+    await arq_redis.enqueue_job('foobar')
+    worker: Worker = worker(functions=[func(foobar, name='foobar')])
+    assert 2 == await worker.run_check()
+
+
+async def test_run_check_error(arq_redis: ArqRedis, worker):
+    await arq_redis.enqueue_job('fails')
+    worker: Worker = worker(functions=[func(fails, name='fails')])
+    with pytest.raises(FailedJobs, match='1 job failed "TypeError: my type error"'):
+        await worker.run_check()
+
+
+async def test_run_check_error2(arq_redis: ArqRedis, worker):
+    await arq_redis.enqueue_job('fails')
+    await arq_redis.enqueue_job('fails')
+    worker: Worker = worker(functions=[func(fails, name='fails')])
+    with pytest.raises(FailedJobs, match='2 jobs failed') as exc_info:
+        await worker.run_check()
+    assert len(exc_info.value.job_results) == 2
+
+
+async def test_return_exception(arq_redis: ArqRedis, worker):
+    async def return_error(ctx):
+        return TypeError('xxx')
+
+    j = await arq_redis.enqueue_job('return_error')
+    worker: Worker = worker(functions=[func(return_error, name='return_error')])
+    await worker.async_run()
+    assert (worker.jobs_complete, worker.jobs_failed, worker.jobs_retried) == (1, 0, 0)
+    r = await j.result(pole_delay=0)
+    assert isinstance(r, TypeError)
+    info = await j.result_info()
+    assert info['success'] is True
+
+
+async def test_error_success(arq_redis: ArqRedis, worker):
+    j = await arq_redis.enqueue_job('fails')
+    worker: Worker = worker(functions=[func(fails, name='fails')])
+    await worker.async_run()
+    assert (worker.jobs_complete, worker.jobs_failed, worker.jobs_retried) == (0, 1, 0)
+    info = await j.result_info()
+    assert info['success'] is False