Fix tests

Author: AlexVelezLl

kolibri/core/tasks/job.py

@@ -4,7 +4,6 @@
 import traceback
 import uuid
 from collections import namedtuple
-from datetime import timedelta
 
 from kolibri.core.tasks.constants import (  # noqa F401 - imported for backwards compatibility
     Priority,
@@ -102,10 +101,6 @@ def default_status_text(job):
 
 ALLOWED_RETRY_IN_KWARGS = {"priority", "repeat", "interval", "retry_interval"}
 
-MAX_RETRIES = (
-    3  # Maximum number of retries for a job that failed due to a retryable exception
-)
-
 
 class Job(object):
     """

kolibri/core/tasks/storage.py

@@ -657,11 +657,9 @@ def reschedule_finished_job_if_needed(  # noqa: C901
             if retry_interval is not NO_VALUE
             else orm_job.retry_interval,
             max_retries=orm_job.max_retries,
+            retries=orm_job.retries,
         )
 
-        if exception is not None:
-            # TODO: Implement retry on exception logic.
-            pass
         # Set a null new_scheduled_time so that we finish processing if none of the cases below pertain.
         new_scheduled_time = None
         if delay is not None:
@@ -671,11 +669,9 @@ def reschedule_finished_job_if_needed(  # noqa: C901
             # enqueuing changes - so if it is still set to repeat, it will repeat again after the
             # delayed rerun.
             new_scheduled_time = self._now() + delay
-        elif self._should_retry_on_task_failed(
+        elif self._should_retry_on_failed_task(
             orm_job, exception, kwargs["retry_interval"]
         ):
-            # If the task has failed, and a retry interval has been specified (either in the original enqueue,
-            # or from the passed in kwargs) then requeue as a retry.
             new_scheduled_time = self._now() + timedelta(
                 seconds=kwargs["retry_interval"]
                 if kwargs["retry_interval"] is not None
@@ -702,7 +698,7 @@ def reschedule_finished_job_if_needed(  # noqa: C901
             # Use the schedule method so that any scheduling hooks are run for this next run of the job.
             self.schedule(new_scheduled_time, job, **kwargs)
 
-    def _should_retry_on_task_failed(self, orm_job, exception, retry_interval):
+    def _should_retry_on_failed_task(self, orm_job, exception, retry_interval):
         """
         Determine if a job should be retried based on its retry settings and the exception raised.
         """
@@ -719,8 +715,8 @@ def _should_retry_on_task_failed(self, orm_job, exception, retry_interval):
 
         job = self._orm_to_job(orm_job)
         retry_on = job.task.retry_on
-        if retry_on:
-            return any(issubclass(exception, exc) for exc in retry_on)
+        if retry_on and exception:
+            return any(isinstance(exception, exc) for exc in retry_on)
 
         return True
 
@@ -831,6 +827,7 @@ def schedule(
         interval=0,
         repeat=0,
         retry_interval=None,
+        retries=None,
         max_retries=None,
     ):
         """
@@ -867,6 +864,7 @@ def schedule(
                 interval=interval,
                 repeat=repeat,
                 retry_interval=retry_interval,
+                retries=retries,
                 max_retries=max_retries,
                 scheduled_time=naive_utc_datetime(dt),
                 saved_job=job.to_json(),

kolibri/core/tasks/test/taskrunner/test_job_running.py

@@ -18,8 +18,9 @@
 @pytest.fixture
 def storage_fixture():
     with connection() as conn:
-        e = Worker(connection=conn)
         b = Storage(conn)
+        b.Base.metadata.create_all(conn)
+        e = Worker(connection=conn)
         b.clear(force=True)
         yield b
         e.shutdown()

kolibri/core/tasks/test/taskrunner/test_scheduler.py

@@ -2,6 +2,7 @@
 
 import pytest
 
+from kolibri.core.tasks.decorators import register_task
 from kolibri.core.tasks.exceptions import JobRunning
 from kolibri.core.tasks.job import Job
 from kolibri.core.tasks.storage import Storage
@@ -14,15 +15,21 @@
 def job_storage():
     with connection() as c:
         s = Storage(connection=c)
+        s.Base.metadata.create_all(c)
         s.clear(force=True)
         yield s
         s.clear(force=True)
 
 
+@register_task
+def add(x, y):
+    return x + y
+
+
 class TestScheduler(object):
     @pytest.fixture
     def job(self):
-        return Job(id)
+        return Job(add)
 
     def test_enqueue_at_a_function(self, job_storage, job):
         job_id = job_storage.enqueue_at(local_now(), job)
@@ -171,7 +178,7 @@ class TestReschedule(TestScheduler):
     @pytest.fixture
     def job(self, job_storage):
         now = local_now()
-        job_id = job_storage.schedule(now, Job(id), interval=1, repeat=123)
+        job_id = job_storage.schedule(now, Job(add), interval=1, repeat=123)
         return job_storage.get_job(job_id)
 
     def test_reschedule_a_function_gives_job_running_error(self, job_storage, job):

kolibri/core/tasks/test/taskrunner/test_storage.py

@@ -5,6 +5,7 @@
 import pytest
 import pytz
 from mock import patch
+from requests.exceptions import HTTPError
 
 from kolibri.core.tasks.constants import DEFAULT_QUEUE
 from kolibri.core.tasks.constants import Priority
@@ -26,36 +27,43 @@
 def defaultbackend():
     with connection() as c:
         b = Storage(c)
+        b.Base.metadata.create_all(c)
         b.clear(force=True)
         yield b
         b.clear(force=True)
 
 
-@pytest.fixture
-def func():
-    @register_task
-    def add(x, y):
-        return x + y
+@register_task(
+    retry_on=[ValueError, TypeError],
+)
+def add(x, y):
+    return x + y
 
-    TaskRegistry["kolibri.core.tasks.test.taskrunner.test_storage.add"] = add
 
-    yield add
-    TaskRegistry.clear()
+@pytest.fixture(autouse=True)
+def register_add_task():
+    # register before tests
+    TaskRegistry[callable_to_import_path(add)] = add
+    try:
+        yield
+    finally:
+        # clear after tests
+        TaskRegistry.clear()
 
 
 @pytest.fixture
-def simplejob(func):
-    return Job(func)
+def simplejob():
+    return Job(add)
 
 
 class TestBackend:
-    def test_can_enqueue_single_job(self, defaultbackend, simplejob, func):
+    def test_can_enqueue_single_job(self, defaultbackend, simplejob):
         job_id = defaultbackend.enqueue_job(simplejob, QUEUE)
 
         new_job = defaultbackend.get_job(job_id)
 
         # Does the returned job record the function we set to run?
-        assert str(new_job.func) == callable_to_import_path(func)
+        assert str(new_job.func) == callable_to_import_path(add)
 
         # Does the job have the right state (QUEUED)?
         assert new_job.state == State.QUEUED
@@ -310,6 +318,134 @@ def test_can_reschedule_finished_job(self, defaultbackend, simplejob):
         assert requeued_job.state == State.QUEUED
         assert requeued_orm_job.scheduled_time > previous_scheduled_time
 
+    def test_job_retry_on_matching_exception(self, defaultbackend, simplejob):
+        exception = ValueError("Error")
+        job_id = defaultbackend.enqueue_job(
+            simplejob, QUEUE, retry_interval=5, max_retries=3
+        )
+        defaultbackend.mark_job_as_failed(job_id, exception, "Traceback")
+
+        orm_job = defaultbackend.get_orm_job(job_id)
+        previous_scheduled_time = orm_job.scheduled_time
+
+        defaultbackend.reschedule_finished_job_if_needed(
+            simplejob.job_id, exception=exception
+        )
+        requeued_orm_job = defaultbackend.get_orm_job(job_id)
+        requeued_job = defaultbackend.get_job(job_id)
+
+        assert requeued_job.state == State.QUEUED
+        assert requeued_orm_job.scheduled_time > previous_scheduled_time
+        assert requeued_orm_job.retries == 1
+
+    def test_job_retry_on_matching_exception__no_max_retries(
+        self, defaultbackend, simplejob
+    ):
+        exception = ValueError("Error")
+        job_id = defaultbackend.enqueue_job(simplejob, QUEUE, retry_interval=5)
+        defaultbackend.mark_job_as_failed(job_id, exception, "Traceback")
+
+        orm_job = defaultbackend.get_orm_job(job_id)
+        previous_scheduled_time = orm_job.scheduled_time
+
+        defaultbackend.reschedule_finished_job_if_needed(
+            simplejob.job_id, exception=exception
+        )
+        requeued_orm_job = defaultbackend.get_orm_job(job_id)
+        requeued_job = defaultbackend.get_job(job_id)
+
+        assert requeued_job.state == State.QUEUED
+        assert requeued_orm_job.scheduled_time > previous_scheduled_time
+        assert requeued_orm_job.retries == 1
+
+    def test_job_retry_on_matching_exception__no_retry_interval(
+        self, defaultbackend, simplejob
+    ):
+        exception = TypeError("Error")
+        job_id = defaultbackend.enqueue_job(simplejob, QUEUE, max_retries=3)
+        defaultbackend.mark_job_as_failed(job_id, exception, "Traceback")
+
+        orm_job = defaultbackend.get_orm_job(job_id)
+        previous_scheduled_time = orm_job.scheduled_time
+
+        defaultbackend.reschedule_finished_job_if_needed(
+            simplejob.job_id, exception=exception
+        )
+        requeued_orm_job = defaultbackend.get_orm_job(job_id)
+        requeued_job = defaultbackend.get_job(job_id)
+
+        assert requeued_job.state == State.QUEUED
+        assert requeued_orm_job.scheduled_time > previous_scheduled_time
+        assert requeued_orm_job.retries == 1
+
+    def test_job_not_retry_on_matching_exception__no_retry_params(
+        self, defaultbackend, simplejob
+    ):
+        # If job has no retry params, it should not retry even if exception matches
+        exception = ValueError("Error")
+        job_id = defaultbackend.enqueue_job(simplejob, QUEUE)
+        defaultbackend.mark_job_as_failed(job_id, exception, "Traceback")
+
+        orm_job = defaultbackend.get_orm_job(job_id)
+        previous_scheduled_time = orm_job.scheduled_time
+
+        defaultbackend.reschedule_finished_job_if_needed(
+            simplejob.job_id, exception=exception
+        )
+        requeued_orm_job = defaultbackend.get_orm_job(job_id)
+        requeued_job = defaultbackend.get_job(job_id)
+
+        assert requeued_job.state == State.FAILED
+        assert requeued_orm_job.scheduled_time == previous_scheduled_time
+        assert requeued_orm_job.retries is None
+
+    def test_job_not_retry_on_non_matching_exception(self, defaultbackend, simplejob):
+        exception = HTTPError("Error")
+        job_id = defaultbackend.enqueue_job(
+            simplejob, QUEUE, retry_interval=5, max_retries=3
+        )
+        defaultbackend.mark_job_as_failed(job_id, exception, "Traceback")
+
+        orm_job = defaultbackend.get_orm_job(job_id)
+        previous_scheduled_time = orm_job.scheduled_time
+
+        defaultbackend.reschedule_finished_job_if_needed(
+            simplejob.job_id, exception=exception
+        )
+        requeued_orm_job = defaultbackend.get_orm_job(job_id)
+        requeued_job = defaultbackend.get_job(job_id)
+
+        assert requeued_job.state == State.FAILED
+        assert requeued_orm_job.scheduled_time == previous_scheduled_time
+        assert requeued_orm_job.retries is None
+
+    def test_job_not_retry_on_limit_max_retries(self, defaultbackend, simplejob):
+        exception = ValueError("Error")
+        job_id = defaultbackend.enqueue_job(simplejob, QUEUE, max_retries=1)
+        defaultbackend.mark_job_as_failed(job_id, exception, "Traceback")
+
+        # Retry first time
+        defaultbackend.reschedule_finished_job_if_needed(
+            simplejob.job_id, exception=exception
+        )
+        defaultbackend.mark_job_as_failed(job_id, exception, "Traceback")
+
+        orm_job = defaultbackend.get_orm_job(job_id)
+        previous_scheduled_time = orm_job.scheduled_time
+
+        # When trying to retry second time, it should not retry as max_retries is reached
+        defaultbackend.reschedule_finished_job_if_needed(
+            simplejob.job_id, exception=exception
+        )
+
+        requeued_orm_job = defaultbackend.get_orm_job(job_id)
+        requeued_job = defaultbackend.get_job(job_id)
+
+        retries = requeued_orm_job.retries
+        assert requeued_job.state == State.FAILED
+        assert requeued_orm_job.scheduled_time == previous_scheduled_time
+        assert retries == 1
+
     def test_reschedule_finished_job_canceled(self, defaultbackend, simplejob):
         # Test case where the job is canceled.
         job_id = defaultbackend.enqueue_job(simplejob, QUEUE)

kolibri/core/tasks/test/taskrunner/test_worker.py

@@ -6,6 +6,7 @@
 from kolibri.core.tasks.constants import Priority
 from kolibri.core.tasks.job import Job
 from kolibri.core.tasks.job import State
+from kolibri.core.tasks.storage import Storage
 from kolibri.core.tasks.test.base import connection
 from kolibri.core.tasks.test.taskrunner.test_job_running import EventProxy
 from kolibri.core.tasks.worker import Worker
@@ -43,16 +44,20 @@ def toggle_flag(flag_id):
 @pytest.fixture
 def worker():
     with connection() as c:
+        s = Storage(c)
+        s.Base.metadata.create_all(c)
         b = Worker(c, regular_workers=1, high_workers=1)
         b.storage.clear(force=True)
         yield b
         b.storage.clear(force=True)
+        s.clear(force=True)
         b.shutdown()
 
 
 def test_keyerror_prevention(worker):
     # Create a job with the same ID as the one in worker.enqueue_job_runs_job
     job = Job(id, args=(9,))
+    worker.storage.Base.metadata.create_all(worker.storage.engine)
     worker.storage.enqueue_job(job, QUEUE)
 
     while job.state != "COMPLETED":