test passing

Author: bisgaard-itis

services/api-server/src/simcore_service_api_server/api/routes/tasks.py

@@ -1,7 +1,8 @@
 import logging
 from typing import Annotated, Any
 
-from fastapi import APIRouter, Depends, FastAPI, status
+from celery.exceptions import CeleryError  # type: ignore[import-untyped]
+from fastapi import APIRouter, Depends, FastAPI, HTTPException, status
 from models_library.api_schemas_long_running_tasks.base import TaskProgress
 from models_library.api_schemas_long_running_tasks.tasks import (
     TaskGet,
@@ -14,14 +15,13 @@
 )
 from models_library.products import ProductName
 from models_library.users import UserID
+from servicelib.celery.models import TaskFilter, TaskUUID
 from servicelib.fastapi.dependencies import get_app
 
 from ...models.schemas.base import ApiServerEnvelope
 from ...models.schemas.errors import ErrorGet
-from ...services_rpc.async_jobs import AsyncJobClient
 from ..dependencies.authentication import get_current_user_id, get_product_name
-from ..dependencies.celery import ASYNC_JOB_CLIENT_NAME
-from ..dependencies.tasks import get_async_jobs_client
+from ..dependencies.celery import ASYNC_JOB_CLIENT_NAME, get_task_manager_from_app
 from ._constants import (
     FMSG_CHANGELOG_NEW_IN_VERSION,
     create_route_description,
@@ -31,10 +31,11 @@
 _logger = logging.getLogger(__name__)
 
 
-def _get_job_filter(user_id: UserID, product_name: ProductName) -> AsyncJobFilter:
-    return AsyncJobFilter(
+def _get_task_filter(user_id: UserID, product_name: ProductName) -> TaskFilter:
+    job_filter = AsyncJobFilter(
         user_id=user_id, product_name=product_name, client_name=ASYNC_JOB_CLIENT_NAME
     )
+    return TaskFilter.model_validate(job_filter.model_dump())
 
 
 _DEFAULT_TASK_STATUS_CODES: dict[int | str, dict[str, Any]] = {
@@ -61,26 +62,34 @@ async def list_tasks(
     app: Annotated[FastAPI, Depends(get_app)],
     user_id: Annotated[UserID, Depends(get_current_user_id)],
     product_name: Annotated[ProductName, Depends(get_product_name)],
-    async_jobs: Annotated[AsyncJobClient, Depends(get_async_jobs_client)],
 ):
-    user_async_jobs = await async_jobs.list_jobs(
-        job_filter=_get_job_filter(user_id, product_name),
-        filter_="",
-    )
+
+    task_manager = get_task_manager_from_app(app)
+
+    try:
+        tasks = await task_manager.list_tasks(
+            task_filter=_get_task_filter(user_id, product_name),
+        )
+    except CeleryError as exc:
+        raise HTTPException(
+            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+            detail="Encountered issue when listing tasks",
+        ) from exc
+
     app_router = app.router
     data = [
         TaskGet(
-            task_id=f"{job.job_id}",
-            task_name=job.job_name,
+            task_id=f"{task.uuid}",
+            task_name=task.metadata.name,
             status_href=app_router.url_path_for(
-                "get_task_status", task_id=f"{job.job_id}"
+                "get_task_status", task_id=f"{task.uuid}"
             ),
-            abort_href=app_router.url_path_for("cancel_task", task_id=f"{job.job_id}"),
+            abort_href=app_router.url_path_for("cancel_task", task_id=f"{task.uuid}"),
             result_href=app_router.url_path_for(
-                "get_task_result", task_id=f"{job.job_id}"
+                "get_task_result", task_id=f"{task.uuid}"
             ),
         )
-        for job in user_async_jobs
+        for task in tasks
     ]
     return ApiServerEnvelope(data=data)
 
@@ -99,20 +108,29 @@ async def list_tasks(
 )
 async def get_task_status(
     task_id: AsyncJobId,
+    app: Annotated[FastAPI, Depends(get_app)],
     user_id: Annotated[UserID, Depends(get_current_user_id)],
     product_name: Annotated[ProductName, Depends(get_product_name)],
-    async_jobs: Annotated[AsyncJobClient, Depends(get_async_jobs_client)],
 ):
-    async_job_rpc_status = await async_jobs.status(
-        job_id=task_id,
-        job_filter=_get_job_filter(user_id, product_name),
-    )
-    _task_id = f"{async_job_rpc_status.job_id}"
+    task_manager = get_task_manager_from_app(app)
+
+    try:
+        task_status = await task_manager.get_task_status(
+            task_filter=_get_task_filter(user_id, product_name),
+            task_uuid=TaskUUID(f"{task_id}"),
+        )
+    except CeleryError as exc:
+        raise HTTPException(
+            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+            detail="Encountered issue when getting task status",
+        ) from exc
+
     return TaskStatus(
         task_progress=TaskProgress(
-            task_id=_task_id, percent=async_job_rpc_status.progress.percent_value
+            task_id=f"{task_status.task_uuid}",
+            percent=task_status.progress_report.percent_value,
         ),
-        done=async_job_rpc_status.done,
+        done=task_status.is_done,
         started=None,
     )
 
@@ -131,14 +149,22 @@ async def get_task_status(
 )
 async def cancel_task(
     task_id: AsyncJobId,
+    app: Annotated[FastAPI, Depends(get_app)],
     user_id: Annotated[UserID, Depends(get_current_user_id)],
     product_name: Annotated[ProductName, Depends(get_product_name)],
-    async_jobs: Annotated[AsyncJobClient, Depends(get_async_jobs_client)],
 ):
-    await async_jobs.cancel(
-        job_id=task_id,
-        job_filter=_get_job_filter(user_id, product_name),
-    )
+    task_manager = get_task_manager_from_app(app)
+
+    try:
+        await task_manager.cancel_task(
+            task_filter=_get_task_filter(user_id, product_name),
+            task_uuid=TaskUUID(f"{task_id}"),
+        )
+    except CeleryError as exc:
+        raise HTTPException(
+            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+            detail="Encountered issue when cancelling task",
+        ) from exc
 
 
 @router.get(
@@ -165,12 +191,34 @@ async def cancel_task(
 )
 async def get_task_result(
     task_id: AsyncJobId,
+    app: Annotated[FastAPI, Depends(get_app)],
     user_id: Annotated[UserID, Depends(get_current_user_id)],
     product_name: Annotated[ProductName, Depends(get_product_name)],
-    async_jobs: Annotated[AsyncJobClient, Depends(get_async_jobs_client)],
 ):
-    async_job_rpc_result = await async_jobs.result(
-        job_id=task_id,
-        job_filter=_get_job_filter(user_id, product_name),
-    )
-    return TaskResult(result=async_job_rpc_result.result, error=None)
+    task_manager = get_task_manager_from_app(app)
+    task_filter = _get_task_filter(user_id, product_name)
+
+    try:
+        # First check if task exists and is done
+        task_status = await task_manager.get_task_status(
+            task_filter=task_filter,
+            task_uuid=TaskUUID(f"{task_id}"),
+        )
+
+        if not task_status.is_done:
+            raise HTTPException(
+                status_code=status.HTTP_404_NOT_FOUND,
+                detail="Task result not available yet",
+            )
+
+        result = await task_manager.get_task_result(
+            task_filter=task_filter,
+            task_uuid=TaskUUID(f"{task_id}"),
+        )
+        return TaskResult(result=result, error=None)
+
+    except CeleryError as exc:
+        raise HTTPException(
+            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+            detail="Encountered issue when getting task result",
+        ) from exc

services/api-server/tests/unit/celery/test_functions.py

@@ -5,6 +5,7 @@
 from celery import Celery, Task
 from celery.contrib.testing.worker import TestWorkController
 from celery_library.task import register_task
+from celery_library.types import register_pydantic_types
 from faker import Faker
 from fastapi import FastAPI, status
 from httpx import AsyncClient, BasicAuth
@@ -111,6 +112,7 @@ async def run_function(
     ), f"Signature mismatch: {inspect.signature(run_function_task)} != {inspect.signature(run_function)}"
 
     def _(celery_app: Celery) -> None:
+        register_pydantic_types(RegisteredProjectFunctionJob)
         register_task(celery_app, run_function)
 
     return _
@@ -148,8 +150,4 @@ async def test_with_fake_run_function(
 
     # Poll until task completion and get result
     result = await poll_task_until_done(client, auth, task.task_id)
-
-    # Verify the result is a RegisteredProjectFunctionJob
-    assert result is not None
-    assert isinstance(result, dict)
-    # Add more specific assertions based on your expected result structure
+    RegisteredProjectFunctionJob.model_validate(result.result)