continue

Author: giancarloromeo

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

@@ -26,18 +26,20 @@
 _STARTUP_TIMEOUT: Final[float] = datetime.timedelta(minutes=1).total_seconds()
 
 
-def on_worker_init(sender, **_kwargs) -> None:
+def on_worker_init(
+    app_factory,
+    celery_settings,
+    sender,
+    **_kwargs,
+) -> None:
     startup_complete_event = threading.Event()
 
     def _init(startup_complete_event: threading.Event) -> None:
         loop = asyncio.new_event_loop()
         asyncio.set_event_loop(loop)
         shutdown_event = asyncio.Event()
 
-        fastapi_app = create_app(app_settings)
-
-        assert app_settings.STORAGE_CELERY
-        celery_settings = app_settings.STORAGE_CELERY
+        fastapi_app = app_factory()
 
         async def setup_task_worker():
             redis_client_sdk = RedisClientSDK(

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

@@ -1,6 +1,7 @@
 """Main application to be deployed in for example uvicorn."""
 
 import logging
+from functools import partial
 
 from celery.signals import worker_init, worker_shutdown  # type: ignore[import-untyped]
 from celery_library.common import create_app as create_celery_app
@@ -11,6 +12,7 @@
 from servicelib.logging_utils import config_all_loggers
 from simcore_service_storage.api._worker_tasks.tasks import setup_worker_tasks
 
+from ...core.application import create_app
 from ...core.settings import ApplicationSettings
 
 _settings = ApplicationSettings.create_from_envs()
@@ -26,7 +28,8 @@
 
 assert _settings.STORAGE_CELERY
 app = create_celery_app(_settings.STORAGE_CELERY)
-worker_init.connect(on_worker_init)
+app_factory = partial(create_app(_settings))
+worker_init.connect(partial(on_worker_init, app_factory, _settings.STORAGE_CELERY))
 worker_shutdown.connect(on_worker_shutdown)
 
 

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

@@ -12,6 +12,7 @@
 import random
 import sys
 from collections.abc import AsyncIterator, Awaitable, Callable
+from functools import partial
 from pathlib import Path
 from typing import Any, Final, cast
 
@@ -974,10 +975,7 @@ def celery_config() -> dict[str, Any]:
 def mock_celery_app(mocker: MockerFixture, celery_config: dict[str, Any]) -> Celery:
     celery_app = Celery(**celery_config)
 
-    for module in (
-        "simcore_service_storage.modules.celery._common.create_app",
-        "simcore_service_storage.modules.celery.create_app",
-    ):
+    for module in ("celery_library.create_app",):
         mocker.patch(module, return_value=celery_app)
 
     return celery_app
@@ -1000,13 +998,17 @@ async def with_storage_celery_worker_controller(
     register_celery_tasks: Callable[[Celery], None],
 ) -> AsyncIterator[TestWorkController]:
     # Signals must be explicitily connected
-    worker_init.connect(on_worker_init)
+    monkeypatch.setenv("STORAGE_WORKER_MODE", "true")
+    app_settings = ApplicationSettings.create_from_envs()
+    app_factory = partial(create_app, app_settings)
+    worker_init.connect(
+        partial(on_worker_init, app_factory, app_settings.STORAGE_CELERY)
+    )
     worker_shutdown.connect(on_worker_shutdown)
 
     setup_worker_tasks(celery_app)
     register_celery_tasks(celery_app)
 
-    monkeypatch.setenv("STORAGE_WORKER_MODE", "true")
     with start_worker(
         celery_app,
         pool="threads",