diff --git a/backoffice/backoffice/workflows/airflow_utils.py b/backoffice/backoffice/workflows/airflow_utils.py index 45a5e8c9..914b7816 100644 --- 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 index 40641cea..734af996 100644 --- 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 index c6863a47..273faeac 100644 --- 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 index 60c928fd..ecb69660 100644 --- a/backoffice/backoffice/workflows/tests/test_airflow_utils.py +++ b/backoffice/backoffice/workflows/tests/test_airflow_utils.py @@ -15,7 +15,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) ) diff --git a/backoffice/backoffice/workflows/tests/test_views.py b/backoffice/backoffice/workflows/tests/test_views.py index fc40b5df..1a733829 100644 --- a/backoffice/backoffice/workflows/tests/test_views.py +++ b/backoffice/backoffice/workflows/tests/test_views.py @@ -261,7 +261,7 @@ 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, ) @@ -269,7 +269,7 @@ def setUp(self): 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()