feat(ws): add type and study_id inside TaskEventPayload

AntaresSimulatorTeam · Oct 30, 2024 · 04bd2ca · 04bd2ca
1 parent c66b419
commit 04bd2ca
Show file tree

Hide file tree

Showing 4 changed files with 17 additions and 5 deletions.
diff --git a/antarest/core/tasks/model.py b/antarest/core/tasks/model.py
@@ -79,7 +79,7 @@ class TaskEventPayload(AntaresBaseModel, extra="forbid"):
     id: str
     message: str
     type: TaskType
-    study_id: str
+    study_id: t.Optional[str] = None
 
 
 class TaskDTO(AntaresBaseModel, extra="forbid"):

diff --git a/antarest/core/tasks/service.py b/antarest/core/tasks/service.py
@@ -289,6 +289,8 @@ def _launch_task(
                     message=custom_event_messages.start
                     if custom_event_messages is not None
                     else f"Task {task.id} added",
+                    type=task.type,
+                    study_id=task.ref_id,
                 ).model_dump(),
                 permissions=PermissionInfo(owner=request_params.user.impersonator),
             )
@@ -389,6 +391,10 @@ def _run_task(
         custom_event_messages: t.Optional[CustomTaskEventMessages] = None,
     ) -> None:
         # attention: this function is executed in a thread, not in the main process
+        with db():
+            task = db.session.query(TaskJob).get(task_id)
+            task_type = task.type
+            study_id = task.ref_id
 
         self.event_bus.push(
             Event(
@@ -398,6 +404,8 @@ def _run_task(
                     message=custom_event_messages.running
                     if custom_event_messages is not None
                     else f"Task {task_id} is running",
+                    type=task_type,
+                    study_id=study_id,
                 ).model_dump(),
                 permissions=PermissionInfo(public_mode=PublicMode.READ),
                 channel=EventChannelDirectory.TASK + task_id,
@@ -444,6 +452,8 @@ def _run_task(
                             if custom_event_messages is not None
                             else f"Task {task_id} {event_msg}"
                         ),
+                        type=task_type,
+                        study_id=study_id,
                     ).model_dump(),
                     permissions=PermissionInfo(public_mode=PublicMode.READ),
                     channel=EventChannelDirectory.TASK + task_id,
@@ -469,7 +479,9 @@ def _run_task(
             self.event_bus.push(
                 Event(
                     type=EventType.TASK_FAILED,
-                    payload=TaskEventPayload(id=task_id, message=message).model_dump(),
+                    payload=TaskEventPayload(
+                        id=task_id, message=message, type=task_type, study_id=study_id
+                    ).model_dump(),
                     permissions=PermissionInfo(public_mode=PublicMode.READ),
                     channel=EventChannelDirectory.TASK + task_id,
                 )

diff --git a/tests/core/test_tasks.py b/tests/core/test_tasks.py
@@ -143,7 +143,7 @@ def action_fail(notifier: ITaskNotifier) -> TaskResult:
     failed_id = service.add_task(
         action_fail,
         "failed action",
-        None,
+        TaskType.COPY,
         None,
         None,
         None,
@@ -172,7 +172,7 @@ def action_ok(notifier: ITaskNotifier) -> TaskResult:
     ok_id = service.add_task(
         action_ok,
         None,
-        None,
+        TaskType.COPY,
         None,
         None,
         None,

diff --git a/webapp/src/services/webSocket/types.ts b/webapp/src/services/webSocket/types.ts
@@ -60,7 +60,7 @@ export interface TaskEventPayload {
   id: string;
   message: string;
   type: TTaskType;
-  study_id: StudyMetadata["id"];
+  study_id?: StudyMetadata["id"];
 }
 
 export interface TsGenerationProgressPayload {