fix tests

Author: giancarloromeo

packages/celery-library/src/celery_library/client.py

@@ -54,7 +54,7 @@ async def submit_task(
             self._celery_app.send_task(
                 task_metadata.name,
                 task_id=task_id,
-                kwargs={"task_id": task_id} | task_params,
+                kwargs=task_params,
                 queue=task_metadata.queue.value,
             )
 

packages/celery-library/src/celery_library/task.py

@@ -40,11 +40,11 @@ class TaskAbortedError(Exception): ...
 def _async_task_wrapper(
     app: Celery,
 ) -> Callable[
-    [Callable[Concatenate[AbortableTask, P], Coroutine[Any, Any, R]]],
+    [Callable[Concatenate[AbortableTask, TaskId, P], Coroutine[Any, Any, R]]],
     Callable[Concatenate[AbortableTask, P], R],
 ]:
     def decorator(
-        coro: Callable[Concatenate[AbortableTask, P], Coroutine[Any, Any, R]],
+        coro: Callable[Concatenate[AbortableTask, TaskId, P], Coroutine[Any, Any, R]],
     ) -> Callable[Concatenate[AbortableTask, P], R]:
         @wraps(coro)
         def wrapper(task: AbortableTask, *args: P.args, **kwargs: P.kwargs) -> R:
@@ -56,7 +56,7 @@ async def run_task(task_id: TaskID) -> R:
                 try:
                     async with asyncio.TaskGroup() as tg:
                         main_task = tg.create_task(
-                            coro(task, *args, **kwargs),
+                            coro(task, task_id, *args, **kwargs),
                         )
 
                         async def abort_monitor():
@@ -205,5 +205,4 @@ def register_task(  # type: ignore[misc]
         bind=True,
         base=AbortableTask,
         time_limit=None if timeout is None else timeout.total_seconds(),
-        pydantic=True,
     )(wrapped_fn)

packages/celery-library/src/celery_library/types.py

@@ -1,6 +1,9 @@
+from functools import partial
 from pathlib import Path
+from typing import Any
 
-from kombu.utils.json import register_type  # type: ignore[import-untyped]
+from kombu.utils.json import register_type
+from pydantic import BaseModel  # type: ignore[import-untyped]
 
 
 def _path_encoder(obj):
@@ -20,6 +23,14 @@ def _class_full_name(clz: type) -> str:
     return ".".join([clz.__module__, clz.__qualname__])
 
 
+def _pydantic_model_encoder(obj: BaseModel, *args, **kwargs) -> dict[str, Any]:
+    return obj.model_dump(*args, **kwargs, mode="json")
+
+
+def _pydantic_model_decoder(clz: type[BaseModel], data: dict[str, Any]) -> BaseModel:
+    return clz(**data)
+
+
 def register_celery_types() -> None:
     register_type(
         Path,
@@ -28,3 +39,13 @@ def register_celery_types() -> None:
         _path_decoder,
     )
     register_type(set, _class_full_name(set), encoder=list, decoder=set)
+
+
+def register_pydantic_types(*models: type[BaseModel]) -> None:
+    for model in models:
+        register_type(
+            model,
+            _class_full_name(model),
+            encoder=_pydantic_model_encoder,
+            decoder=partial(_pydantic_model_decoder, model),
+        )

services/storage/src/simcore_service_storage/api/_worker_tasks/tasks.py

@@ -2,10 +2,15 @@
 
 from celery import Celery  # type: ignore[import-untyped]
 from celery_library.task import register_task
-from celery_library.types import register_celery_types
+from celery_library.types import register_celery_types, register_pydantic_types
 from models_library.api_schemas_storage.export_data_async_jobs import AccessRightError
+from models_library.api_schemas_storage.storage_schemas import (
+    FileUploadCompletionBody,
+    FoldersBody,
+)
 from servicelib.logging_utils import log_context
 
+from ...models import FileMetaData
 from ._files import complete_upload_file
 from ._paths import compute_path_size, delete_paths
 from ._simcore_s3 import deep_copy_files_from_project, export_data
@@ -16,6 +21,8 @@
 def setup_worker_tasks(app: Celery) -> None:
     register_celery_types()
 
+    register_pydantic_types(FileUploadCompletionBody, FileMetaData, FoldersBody)
+
     with log_context(_logger, logging.INFO, msg="worker task registration"):
         register_task(app, export_data, dont_autoretry_for=(AccessRightError,))
         register_task(app, compute_path_size)