fix tests

Author: giancarloromeo

services/storage/src/simcore_service_storage/modules/celery/signals.py

@@ -0,0 +1,70 @@
+import asyncio
+import logging
+import threading
+from typing import Final
+
+from asgi_lifespan import LifespanManager
+from celery import Celery
+from fastapi import FastAPI
+from servicelib.async_utils import cancel_wait_task
+from simcore_service_storage.core.application import create_app
+from simcore_service_storage.core.settings import ApplicationSettings
+from simcore_service_storage.modules.celery import get_event_loop, set_event_loop
+from simcore_service_storage.modules.celery.utils import (
+    get_fastapi_app,
+    set_celery_worker,
+    set_fastapi_app,
+)
+from simcore_service_storage.modules.celery.worker import CeleryTaskQueueWorker
+
+_logger = logging.getLogger(__name__)
+
+_LIFESPAN_TIMEOUT: Final[int] = 10
+
+
+def on_worker_init(sender, **_kwargs):
+    def _init_fastapi():
+        loop = asyncio.new_event_loop()
+        asyncio.set_event_loop(loop)
+        shutdown_event = asyncio.Event()
+
+        fastapi_app = create_app(ApplicationSettings.create_from_envs())
+
+        async def lifespan():
+            async with LifespanManager(
+                fastapi_app,
+                startup_timeout=_LIFESPAN_TIMEOUT,
+                shutdown_timeout=_LIFESPAN_TIMEOUT,
+            ):
+                try:
+                    await shutdown_event.wait()
+                except asyncio.CancelledError:
+                    _logger.warning("Lifespan task cancelled")
+
+        lifespan_task = loop.create_task(lifespan())
+        fastapi_app.state.lifespan_task = lifespan_task
+        fastapi_app.state.shutdown_event = shutdown_event
+        set_event_loop(fastapi_app, loop)
+
+        set_fastapi_app(sender.app, fastapi_app)
+        set_celery_worker(sender.app, CeleryTaskQueueWorker(sender.app))
+
+        loop.run_forever()
+
+    thread = threading.Thread(target=_init_fastapi, daemon=True)
+    thread.start()
+
+
+def on_worker_shutdown(sender, **_kwargs):
+    assert isinstance(sender.app, Celery)
+
+    fastapi_app = get_fastapi_app(sender.app)
+    assert isinstance(fastapi_app, FastAPI)
+    event_loop = get_event_loop(fastapi_app)
+
+    async def shutdown():
+        fastapi_app.state.shutdown_event.set()
+
+        await cancel_wait_task(fastapi_app.state.lifespan_task, max_delay=5)
+
+    asyncio.run_coroutine_threadsafe(shutdown(), event_loop)

services/storage/src/simcore_service_storage/modules/celery/tasks.py

@@ -0,0 +1,27 @@
+import logging
+import time
+
+from celery import Task
+from models_library.progress_bar import ProgressReport
+from models_library.projects_nodes_io import StorageFileID
+from servicelib.logging_utils import log_context
+from simcore_service_storage.modules.celery.utils import get_celery_worker
+
+_logger = logging.getLogger(__name__)
+
+
+def export_data(task: Task, files: list[StorageFileID]):
+    for n, file in enumerate(files, start=1):
+        with log_context(
+            _logger,
+            logging.INFO,
+            msg=f"Exporting {file=} ({n}/{len(files)})",
+        ):
+            assert task.name
+            get_celery_worker(task.app).set_task_progress(
+                task_name=task.name,
+                task_id=task.request.id,
+                report=ProgressReport(actual_value=n / len(files) * 100),
+            )
+            time.sleep(10)
+    return "done"

services/storage/src/simcore_service_storage/modules/celery/worker_main.py

@@ -1,27 +1,18 @@
 """Main application to be deployed in for example uvicorn."""
 
-import asyncio
 import logging
-import threading
-from typing import Final
 
-from asgi_lifespan import LifespanManager
-from celery import Celery
+from celery.contrib.abortable import AbortableTask
 from celery.signals import worker_init, worker_shutdown
-from fastapi import FastAPI
-from servicelib.background_task import cancel_wait_task
 from servicelib.logging_utils import config_all_loggers
-from simcore_service_storage.modules.celery import get_event_loop, set_event_loop
-from simcore_service_storage.modules.celery.utils import (
-    CeleryTaskQueueWorker,
-    get_fastapi_app,
-    set_celery_worker,
-    set_fastapi_app,
+from simcore_service_storage.modules.celery.signals import (
+    on_worker_init,
+    on_worker_shutdown,
 )
 
-from ...core.application import create_app
 from ...core.settings import ApplicationSettings
 from ._common import create_app as create_celery_app
+from .tasks import export_data
 
 _settings = ApplicationSettings.create_from_envs()
 
@@ -35,58 +26,8 @@
 
 _logger = logging.getLogger(__name__)
 
-_LIFESPAN_TIMEOUT: Final[int] = 10
-
-
-@worker_init.connect
-def on_worker_init(sender, **_kwargs):
-    def _init_fastapi():
-        loop = asyncio.new_event_loop()
-        asyncio.set_event_loop(loop)
-        shutdown_event = asyncio.Event()
-
-        fastapi_app = create_app(_settings)
-
-        async def lifespan():
-            async with LifespanManager(
-                fastapi_app,
-                startup_timeout=_LIFESPAN_TIMEOUT,
-                shutdown_timeout=_LIFESPAN_TIMEOUT,
-            ):
-                try:
-                    await shutdown_event.wait()
-                except asyncio.CancelledError:
-                    _logger.warning("Lifespan task cancelled")
-
-        lifespan_task = loop.create_task(lifespan())
-        fastapi_app.state.lifespan_task = lifespan_task
-        fastapi_app.state.shutdown_event = shutdown_event
-        set_event_loop(fastapi_app, loop)
-
-        set_fastapi_app(sender.app, fastapi_app)
-        set_celery_worker(sender.app, CeleryTaskQueueWorker(sender.app))
-
-        loop.run_forever()
-
-    thread = threading.Thread(target=_init_fastapi, daemon=True)
-    thread.start()
-
-
-@worker_shutdown.connect
-def on_worker_shutdown(sender, **_kwargs):
-    assert isinstance(sender.app, Celery)
-
-    fastapi_app = get_fastapi_app(sender.app)
-    assert isinstance(fastapi_app, FastAPI)
-    event_loop = get_event_loop(fastapi_app)
-
-    async def shutdown():
-        fastapi_app.state.shutdown_event.set()
-
-        await cancel_wait_task(fastapi_app.state.lifespan_task, max_delay=5)
-
-    asyncio.run_coroutine_threadsafe(shutdown(), event_loop)
-
-
 assert _settings.STORAGE_CELERY
 app = create_celery_app(_settings.STORAGE_CELERY)
+worker_init.connect(on_worker_init)
+worker_shutdown.connect(on_worker_shutdown)
+app.task(name="export_data", bind=True, base=AbortableTask)(export_data)