user actions: convert constants to text choices

inspirehep · Jul 30, 2024 · a243824 · a243824
1 parent e45de6c
commit a243824
Show file tree

Hide file tree

Showing 5 changed files with 22 additions and 17 deletions.
diff --git a/backoffice/backoffice/workflows/airflow_utils.py b/backoffice/backoffice/workflows/airflow_utils.py
@@ -159,7 +159,7 @@ def restart_workflow_dags(workflow_id, workflow_type, params=None):
         delete_workflow_dag(dag_id, workflow_id)
 
         return trigger_airflow_dag(
-            WORKFLOW_DAGS[workflow_type][0], str(workflow_id), params
+            WORKFLOW_DAGS[workflow_type].initialize, str(workflow_id), params
         )
 
     return JsonResponse(

diff --git a/backoffice/backoffice/workflows/api/views.py b/backoffice/backoffice/workflows/api/views.py
@@ -105,11 +105,13 @@ def create(self, request):
             )
         logger.info(
             "Trigger Airflow DAG: %s for %s",
-            WORKFLOW_DAGS[workflow.workflow_type][0],
+            WORKFLOW_DAGS[workflow.workflow_type].initialize,
             workflow.id,
         )
         return airflow_utils.trigger_airflow_dag(
-            WORKFLOW_DAGS[workflow.workflow_type][0], str(workflow.id), workflow.data
+            WORKFLOW_DAGS[workflow.workflow_type].initialize,
+            str(workflow.id),
+            workflow.data,
         )
 
     @action(detail=True, methods=["post"])

diff --git a/backoffice/backoffice/workflows/constants.py b/backoffice/backoffice/workflows/constants.py
@@ -35,13 +35,19 @@ class ResolutionDags(models.TextChoices):
     reject = "reject", "author_create_rejected_dag"
 
 
+class AuthorCreateDags(models.TextChoices):
+    initialize = "author_create_initialization_dag", "initialize"
+    approve = "author_create_approved_dag", "approve"
+    reject = "author_create_rejected_dag", "reject"
+
+
+class AuthorUpdateDags(models.TextChoices):
+    initialize = "author_update_dag", "initialize"
+
+
 WORKFLOW_DAGS = {
     WorkflowType.HEP_CREATE: "",
     WorkflowType.HEP_UPDATE: "",
-    WorkflowType.AUTHOR_CREATE: (
-        "author_create_initialization_dag",
-        "author_create_approved_dag",
-        "author_create_rejected_dag",
-    ),
-    WorkflowType.AUTHOR_UPDATE: ("author_update_dag",),
+    WorkflowType.AUTHOR_CREATE: AuthorCreateDags,
+    WorkflowType.AUTHOR_UPDATE: AuthorUpdateDags,
 }
diff --git a/backoffice/backoffice/workflows/tests/test_airflow_utils.py b/backoffice/backoffice/workflows/tests/test_airflow_utils.py
@@ -1,4 +1,3 @@
-import time
 import uuid
 
 import pytest
@@ -15,7 +14,7 @@ class TestAirflowUtils(TransactionTestCase):
     def setUp(self):
         self.workflow_id = uuid.UUID(int=1)
         self.workflow_type = WorkflowType.AUTHOR_CREATE
-        self.dag_id = WORKFLOW_DAGS[self.workflow_type][0]
+        self.dag_id = WORKFLOW_DAGS[self.workflow_type].initialize
         self.response = airflow_utils.trigger_airflow_dag(
             self.dag_id, str(self.workflow_id)
         )
@@ -29,7 +28,6 @@ def test_trigger_airflow_dag(self):
 
     @pytest.mark.vcr()
     def test_restart_failed_tasks(self):
-        time.sleep(20)  # wait for dag to fail
         response = airflow_utils.restart_failed_tasks(
             self.workflow_id, self.workflow_type
         )
@@ -45,7 +43,6 @@ def test_find_executed_dags(self):
 
     @pytest.mark.vcr()
     def test_find_failed_dag(self):
-        time.sleep(20)  # wait for dag to fail
         failed_dag = airflow_utils.find_failed_dag(self.workflow_id, self.workflow_type)
         self.assertEqual(self.dag_id, failed_dag)
 

diff --git a/backoffice/backoffice/workflows/tests/test_views.py b/backoffice/backoffice/workflows/tests/test_views.py
@@ -261,15 +261,15 @@ def setUp(self):
             id=uuid.UUID(int=0),
         )
         airflow_utils.trigger_airflow_dag(
-            WORKFLOW_DAGS[self.workflow.workflow_type][0],
+            WORKFLOW_DAGS[self.workflow.workflow_type].initialize,
             str(self.workflow.id),
             self.workflow.data,
         )
 
     def tearDown(self):
         super().tearDown()
         airflow_utils.delete_workflow_dag(
-            WORKFLOW_DAGS[self.workflow.workflow_type][0], str(self.workflow.id)
+            WORKFLOW_DAGS[self.workflow.workflow_type].initialize, str(self.workflow.id)
         )
 
     @pytest.mark.vcr()
@@ -309,7 +309,7 @@ def test_accept_author(self):
         self.assertEqual(response.status_code, 200)
 
         airflow_utils.delete_workflow_dag(
-            WORKFLOW_DAGS["AUTHOR_CREATE"][1], str(self.workflow.id)
+            WORKFLOW_DAGS["AUTHOR_CREATE"].approve, str(self.workflow.id)
         )
 
     @pytest.mark.vcr()
@@ -328,7 +328,7 @@ def test_reject_author(self):
         self.assertEqual(response.status_code, 200)
 
         airflow_utils.delete_workflow_dag(
-            WORKFLOW_DAGS["AUTHOR_CREATE"][2], str(self.workflow.id)
+            WORKFLOW_DAGS["AUTHOR_CREATE"].reject, str(self.workflow.id)
         )
 
     @pytest.mark.vcr()