fixed broken

Author: Andrei Neagu

packages/models-library/src/models_library/api_schemas_long_running_tasks/tasks.py

@@ -2,7 +2,7 @@
 from datetime import datetime
 from typing import Any
 
-from pydantic import BaseModel, field_validator
+from pydantic import BaseModel, model_validator
 
 from .base import TaskId, TaskProgress
 
@@ -20,24 +20,22 @@ class TaskResult(BaseModel):
 
 class TaskBase(BaseModel):
     task_id: TaskId
-
-    # NOTE: task name can always be extraced from the task_id
-    # since it'e encoded inside it (expect when this is ued
-    # with data coming form the celery tasks)
     task_name: str = ""
 
-    @field_validator("task_name", mode="before")
-    @classmethod
-    def populate_task_name_if_not_provided(cls, task_name: str, info):
-        # attempt to extract the task name from the task_id
-        # if this is coming form a long_running_task
-        task_id = info.data["task_id"]
-        if task_id and task_name == "":
-            parts = task_id.split(".")
+    @model_validator(mode="after")
+    def try_populate_task_name_from_task_id(self) -> "TaskBase":
+        # NOTE: currently this model is used to validate tasks coming from
+        # the celery backend and form long_running_tasks
+        # 1. if a task comes from Celery, it will keep it's given name
+        # 2. if a task comes from long_running_tasks, it will extract it form
+        #   the task_id, which looks like "{PREFIX}.{TASK_NAME}.UNIQUE|{UUID}"
+
+        if self.task_id and self.task_name == "":
+            parts = self.task_id.split(".")
             if len(parts) > 1:
-                task_name = urllib.parse.unquote(parts[1])
+                self.task_name = urllib.parse.unquote(parts[1])
 
-        return task_name
+        return self
 
 
 class TaskGet(TaskBase):

packages/models-library/tests/test_api_schemas_long_running_tasks_tasks.py

@@ -0,0 +1,56 @@
+import pytest
+from models_library.api_schemas_long_running_tasks.tasks import TaskGet
+from pydantic import TypeAdapter
+
+
+def _get_data_withut_task_name(task_id: str) -> dict:
+    return {
+        "task_id": task_id,
+        "status_href": "",
+        "result_href": "",
+        "abort_href": "",
+    }
+
+
+@pytest.mark.parametrize(
+    "data, expected_task_name",
+    [
+        (_get_data_withut_task_name("a.b.c.d"), "b"),
+        (_get_data_withut_task_name("a.b.c"), "b"),
+        (_get_data_withut_task_name("a.b"), "b"),
+        (_get_data_withut_task_name("a"), ""),
+    ],
+)
+def test_try_extract_task_name(data: dict, expected_task_name: str) -> None:
+    task_get = TaskGet(**data)
+    assert task_get.task_name == expected_task_name
+
+    task_get = TypeAdapter(TaskGet).validate_python(data)
+    assert task_get.task_name == expected_task_name
+
+
+def _get_data_with_task_name(task_id: str, task_name: str) -> dict:
+    return {
+        "task_id": task_id,
+        "task_name": task_name,
+        "status_href": "",
+        "result_href": "",
+        "abort_href": "",
+    }
+
+
+@pytest.mark.parametrize(
+    "data, expected_task_name",
+    [
+        (_get_data_with_task_name("a.b.c.d", "a_name"), "a_name"),
+        (_get_data_with_task_name("a.b.c", "a_name"), "a_name"),
+        (_get_data_with_task_name("a.b", "a_name"), "a_name"),
+        (_get_data_with_task_name("a", "a_name"), "a_name"),
+    ],
+)
+def test_task_name_is_provided(data: dict, expected_task_name: str) -> None:
+    task_get = TaskGet(**data)
+    assert task_get.task_name == expected_task_name
+
+    task_get = TypeAdapter(TaskGet).validate_python(data)
+    assert task_get.task_name == expected_task_name