Merge pull request #1104 from syucream/fix/enum-based-job-status

Make job status enum

Author: hinashi

airone/celery.py

@@ -64,9 +64,9 @@ def celery_task_failure_email(**kwargs):
 @task_failure.connect()
 def celery_task_failure_update_job_status(**kwargs):
     """This event handler is update job status when an exception error in celery."""
-    from job.models import Job
+    from job.models import Job, JobStatus
 
     job_id = kwargs["args"][0]
     job = Job.objects.get(id=job_id)
-    job.status = Job.STATUS["ERROR"]
+    job.status = JobStatus.ERROR.value
     job.save(update_fields=["status"])

airone/lib/http.py

@@ -13,7 +13,7 @@
 from airone.lib.types import AttrTypes, AttrTypeValue
 from entity import models as entity_models
 from entry import models as entry_models
-from job.models import Job, JobOperation
+from job.models import JobOperation, JobStatus
 from user.models import History, User
 
 
@@ -155,7 +155,7 @@ def render(request, template, context={}):
 
     # set constracts for job
     context["JOB"] = {
-        "STATUS": Job.STATUS,
+        "STATUS": {s.name: s.value for s in JobStatus},
         "OPERATION": {
             "CREATE": JobOperation.CREATE_ENTRY.value,
             "EDIT": JobOperation.EDIT_ENTRY.value,

airone/lib/job.py

@@ -1,28 +1,33 @@
 import functools
+from typing import Any, Callable
 
-from job.models import Job
+from job.models import Job, JobStatus
 
 
-def may_schedule_until_job_is_ready(func):
+def may_schedule_until_job_is_ready(
+    func: Callable[[Any, Job], JobStatus | tuple[JobStatus, str] | None],
+):
     @functools.wraps(func)
-    def wrapper(kls, job_id):
+    def wrapper(kls, job_id: int):
         job = Job.objects.get(id=job_id)
         if job.proceed_if_ready():
             # update Job status from PREPARING to PROCEEDING
-            job.update(Job.STATUS["PROCESSING"])
+            job.update(JobStatus.PROCESSING.value)
 
             try:
                 # running Job processing
-                ret: int | tuple = func(kls, job)
+                ret: JobStatus | tuple[int, str] | None = func(kls, job)
             except Exception:
-                ret = Job.STATUS["ERROR"]
+                ret = JobStatus.ERROR
 
             # update Job status after finishing Job processing
-            if isinstance(ret, int):
-                job.update(status=ret)
-            elif isinstance(ret, tuple) and len(ret) == 2:
-                job.update(status=ret[0], text=ret[1])
+            if isinstance(ret, JobStatus):
+                job.update(status=ret.value)
+            elif (
+                isinstance(ret, tuple) and isinstance(ret[0], JobStatus) and isinstance(ret[1], str)
+            ):
+                job.update(status=ret[0].value, text=ret[1])
             elif not job.is_canceled():
-                job.update(Job.STATUS["DONE"])
+                job.update(JobStatus.DONE.value)
 
     return wrapper

airone/tests/test_celery.py

@@ -3,7 +3,7 @@
 
 from airone.celery import celery_task_failure_email, celery_task_failure_update_job_status
 from airone.lib.test import AironeViewTest
-from job.models import Job
+from job.models import Job, JobStatus
 
 
 class CeleryTest(AironeViewTest):
@@ -33,7 +33,7 @@ def test_celery_task_failure_email(self):
 
     def test_celery_task_failure_update_job_status(self):
         user = self.guest_login()
-        job: Job = Job.objects.create(user=user, status=Job.STATUS["PROCESSING"])
+        job: Job = Job.objects.create(user=user, status=JobStatus.PREPARING.value)
         celery_task_failure_update_job_status(
             task_id="test_task_id",
             exception=Exception("Test"),
@@ -43,4 +43,4 @@ def test_celery_task_failure_update_job_status(self):
             einfo=ExceptionInfo,
         )
         job.refresh_from_db()
-        self.assertEqual(job.status, Job.STATUS["ERROR"])
+        self.assertEqual(job.status, JobStatus.ERROR.value)

api_v1/job/views.py

@@ -5,7 +5,7 @@
 from rest_framework.response import Response
 from rest_framework.views import APIView
 
-from job.models import Job, JobOperation
+from job.models import Job, JobOperation, JobStatus
 from job.settings import CONFIG as JOB_CONFIG
 
 
@@ -18,10 +18,10 @@ def get(self, request, format=None):
 
         constant = {
             "status": {
-                "processing": Job.STATUS["PROCESSING"],
-                "done": Job.STATUS["DONE"],
-                "error": Job.STATUS["ERROR"],
-                "timeout": Job.STATUS["TIMEOUT"],
+                "processing": JobStatus.PROCESSING.value,
+                "done": JobStatus.DONE.value,
+                "error": JobStatus.ERROR.value,
+                "timeout": JobStatus.TIMEOUT.value,
             },
             "operation": {
                 "create": JobOperation.CREATE_ENTRY.value,
@@ -72,14 +72,14 @@ def delete(self, request, format=None):
                 "Failed to find Job(id=%s)" % job_id, status=status.HTTP_400_BAD_REQUEST
             )
 
-        if job.status == Job.STATUS["DONE"]:
+        if job.status == JobStatus.DONE.value:
             return Response("Target job has already been done")
 
         if job.operation not in Job.CANCELABLE_OPERATIONS:
             return Response("Target job cannot be canceled", status=status.HTTP_400_BAD_REQUEST)
 
         # update job.status to be canceled
-        job.update(Job.STATUS["CANCELED"])
+        job.update(JobStatus.CANCELED.value)
 
         return Response("Success to cancel job")
 
@@ -93,9 +93,9 @@ def post(self, request, job_id, format=None):
             )
 
         # check job status before starting processing
-        if job.status == Job.STATUS["DONE"]:
+        if job.status == JobStatus.DONE.value:
             return Response("Target job has already been done")
-        elif job.status == Job.STATUS["PROCESSING"]:
+        elif job.status == JobStatus.PROCESSING.value:
             return Response("Target job is under processing", status=status.HTTP_400_BAD_REQUEST)
 
         # check job target status

api_v1/tests/job/test_api.py

@@ -5,7 +5,7 @@
 from airone.lib.types import AttrTypeValue
 from entity.models import Entity, EntityAttr
 from entry.models import Entry
-from job.models import Job, JobOperation
+from job.models import Job, JobOperation, JobStatus
 from job.settings import CONFIG
 
 # constants using this tests
@@ -64,10 +64,10 @@ def test_get_jobs(self):
         self.assertEqual(
             results["constant"]["status"],
             {
-                "processing": Job.STATUS["PROCESSING"],
-                "done": Job.STATUS["DONE"],
-                "error": Job.STATUS["ERROR"],
-                "timeout": Job.STATUS["TIMEOUT"],
+                "processing": JobStatus.PROCESSING.value,
+                "done": JobStatus.DONE.value,
+                "error": JobStatus.ERROR.value,
+                "timeout": JobStatus.TIMEOUT.value,
             },
         )
 
@@ -119,7 +119,7 @@ def test_rerun_jobs(self):
         self.assertEqual(resp.content, b'"Success to run command"')
 
         job = Job.objects.get(id=job.id)
-        self.assertEqual(job.status, Job.STATUS["DONE"])
+        self.assertEqual(job.status, JobStatus.DONE.value)
         self.assertEqual(entry.attrs.count(), 1)
 
         attrv = entry.attrs.first().get_latest_value()
@@ -152,7 +152,7 @@ def test_rerun_jobs(self):
 
         self.assertEqual(resp.status_code, 200)
         self.assertEqual(resp.content, b'"Success to run command"')
-        self.assertEqual(Job.objects.get(id=job.id).status, Job.STATUS["DONE"])
+        self.assertEqual(Job.objects.get(id=job.id).status, JobStatus.DONE.value)
         self.assertEqual(entry.attrs.first().get_latest_value().value, "fuga")
 
         # make and send a job to copy entry
@@ -161,7 +161,7 @@ def test_rerun_jobs(self):
 
         self.assertEqual(resp.status_code, 200)
         self.assertEqual(resp.content, b'"Success to run command"')
-        self.assertEqual(Job.objects.get(id=job.id).status, Job.STATUS["DONE"])
+        self.assertEqual(Job.objects.get(id=job.id).status, JobStatus.DONE.value)
 
         # checks it's success to clone entry
         new_entry = Entry.objects.get(name="new_entry", schema=entity)
@@ -225,7 +225,7 @@ def test_cancel_job(self):
 
         # make a job
         job = Job.new_delete(user, entry)
-        self.assertEqual(job.status, Job.STATUS["PREPARING"])
+        self.assertEqual(job.status, JobStatus.PREPARING.value)
 
         # send request without any parameters
         resp = self.client.delete("/api/v1/job/", json.dumps({}), "application/json")
@@ -246,7 +246,7 @@ def test_cancel_job(self):
 
         # send request with proper parameter
         job = Job.new_create(user, entry)
-        self.assertEqual(job.status, Job.STATUS["PREPARING"])
+        self.assertEqual(job.status, JobStatus.PREPARING.value)
 
         resp = self.client.delete(
             "/api/v1/job/", json.dumps({"job_id": job.id}), "application/json"
@@ -255,7 +255,7 @@ def test_cancel_job(self):
         self.assertEqual(resp.content, b'"Success to cancel job"')
 
         job.refresh_from_db()
-        self.assertEqual(job.status, Job.STATUS["CANCELED"])
+        self.assertEqual(job.status, JobStatus.CANCELED.value)
 
     def test_hidden_jobs_is_not_shown(self):
         user = self.guest_login()

api_v1/tests/test_api.py

@@ -17,7 +17,7 @@
 from entry.models import Attribute, AttributeValue, Entry
 from entry.settings import CONFIG as ENTRY_CONFIG
 from group.models import Group
-from job.models import Job, JobOperation
+from job.models import Job, JobOperation, JobStatus
 from role.models import Role
 from trigger import tasks as trigger_tasks
 from trigger.models import TriggerCondition
@@ -1182,7 +1182,7 @@ def test_notify_event_of_creating_entry_when_create_entry(self):
         # check notification event was invoked
         entry = Entry.objects.get(id=resp.json()["result"])
         job_notify = Job.objects.get(target=entry, operation=JobOperation.NOTIFY_CREATE_ENTRY.value)
-        self.assertEqual(job_notify.status, Job.STATUS["DONE"])
+        self.assertEqual(job_notify.status, JobStatus.DONE.value)
 
     @mock.patch("entry.tasks.notify_entry_update", mock.Mock(return_value=mock.Mock()))
     @mock.patch(
@@ -1234,7 +1234,7 @@ def test_notify_event_of_updating_entry_with_specifying_id(self):
 
         # check creating NOTIFY_UPDATE_ENTRY is created because of changing Entry name
         job_notify = Job.objects.get(target=entry, operation=JobOperation.NOTIFY_UPDATE_ENTRY.value)
-        self.assertEqual(job_notify.status, Job.STATUS["DONE"])
+        self.assertEqual(job_notify.status, JobStatus.DONE.value)
 
     @mock.patch("entry.tasks.delete_entry.delay", mock.Mock(side_effect=tasks.delete_entry))
     @mock.patch("entry.tasks.notify_entry_delete", mock.Mock(return_value=mock.Mock()))
@@ -1260,7 +1260,7 @@ def test_notify_event_of_deleting_entry(self):
         self.assertEqual(resp.status_code, 200)
 
         job_notify = Job.objects.get(target=entry, operation=JobOperation.NOTIFY_DELETE_ENTRY.value)
-        self.assertEqual(job_notify.status, Job.STATUS["DONE"])
+        self.assertEqual(job_notify.status, JobStatus.DONE.value)
 
     def test_update_entry_that_has_deleted_attribute(self):
         """
@@ -1342,7 +1342,7 @@ def test_create_entry_when_trigger_is_set(self):
         # check trigger action was worked properly
         job_query = Job.objects.filter(operation=JobOperation.MAY_INVOKE_TRIGGER.value)
         self.assertEqual(job_query.count(), 1)
-        self.assertEqual(job_query.first().status, Job.STATUS["DONE"])
+        self.assertEqual(job_query.first().status, JobStatus.DONE.value)
 
         # check created Entry has expected value that is set by TriggerAction
         entry = Entry.objects.get(id=resp.json()["result"])
@@ -1392,7 +1392,7 @@ def test_update_entry_when_trigger_is_set(self):
         # check trigger action was worked properly
         job_query = Job.objects.filter(operation=JobOperation.MAY_INVOKE_TRIGGER.value)
         self.assertEqual(job_query.count(), 1)
-        self.assertEqual(job_query.first().status, Job.STATUS["DONE"])
+        self.assertEqual(job_query.first().status, JobStatus.DONE.value)
 
         # check created Entry has expected value that is set by TriggerAction
         self.assertEqual(entry.get_attrv("val").value, "hoge")

dashboard/tasks.py

@@ -10,7 +10,7 @@
 from airone.celery import app
 from airone.lib.types import AttrTypeValue
 from entry.models import Entry
-from job.models import Job
+from job.models import Job, JobStatus
 
 
 def _csv_export(job: Job, values, recv_data: dict, has_referral: bool) -> Optional[io.StringIO]:
@@ -170,7 +170,7 @@ def export_search_result(self, job_id):
         return
 
     # set flag to indicate that this job starts processing
-    job.update(Job.STATUS["PROCESSING"])
+    job.update(JobStatus.PROCESSING.value)
 
     user = job.user
     recv_data = json.loads(job.params)
@@ -202,4 +202,4 @@ def export_search_result(self, job_id):
 
     # update job status and save it except for the case that target job is canceled.
     if not job.is_canceled():
-        job.update(Job.STATUS["DONE"])
+        job.update(JobStatus.DONE.value)

dashboard/tests/test_view.py

@@ -24,7 +24,7 @@
 from entry import tasks as entry_tasks
 from entry.models import Attribute, AttributeValue, Entry
 from group.models import Group
-from job.models import Job, JobOperation
+from job.models import Job, JobOperation, JobStatus
 from role.models import Role
 from user.models import User
 
@@ -273,7 +273,7 @@ def test_show_advanced_search_results(self):
         # check export task is executed
         job = Job.objects.last()
         self.assertEqual(job.operation, JobOperation.EXPORT_SEARCH_RESULT.value)
-        self.assertEqual(job.status, Job.STATUS["DONE"])
+        self.assertEqual(job.status, JobStatus.DONE.value)
         self.assertEqual(json.loads(job.params), export_params)
 
         # check result is set at cache
@@ -1450,7 +1450,7 @@ def test_duplicate_export(self):
         self.assertEqual(resp.status_code, 200)
 
         # When the job is finished, the processing is passed.
-        job.status = Job.STATUS["DONE"]
+        job.status = JobStatus.DONE.value
         job.save(update_fields=["status"])
         resp = self.client.post(
             reverse("dashboard:export_search_result"),

dashboard/views.py

@@ -14,7 +14,7 @@
 from entry.admin import AttrResource, AttrValueResource, EntryResource
 from entry.models import AttributeValue, Entry
 from entry.settings import CONFIG as CONFIG_ENTRY
-from job.models import Job
+from job.models import Job, JobStatus
 
 from .settings import CONFIG
 
@@ -305,7 +305,7 @@ def advanced_search_result(request):
 )
 def export_search_result(request, recv_data):
     # check whether same job is sent
-    job_status_not_finished = [Job.STATUS["PREPARING"], Job.STATUS["PROCESSING"]]
+    job_status_not_finished = [JobStatus.PREPARING.value, JobStatus.PROCESSING.value]
     if (
         Job.get_job_with_params(request.user, recv_data)
         .filter(status__in=job_status_not_finished)