tests: refactor

Author: giancarloromeo

packages/celery-library/src/celery_library/backends/_memory.py

@@ -0,0 +1,71 @@
+from dataclasses import dataclass
+from datetime import timedelta
+
+from models_library.progress_bar import ProgressReport
+
+from ..models import (
+    Task,
+    TaskContext,
+    TaskID,
+    TaskMetadata,
+    TaskUUID,
+    build_task_id_prefix,
+)
+
+
+@dataclass
+class MemoryTaskInfo:
+    metadata: TaskMetadata
+    progress: ProgressReport
+
+
+class MemoryTaskInfoStore:
+    def __init__(self) -> None:
+        self._tasks: dict[TaskID, MemoryTaskInfo] = {}
+
+    async def create_task(
+        self,
+        task_id: TaskID,
+        task_metadata: TaskMetadata,
+        expiry: timedelta,
+    ) -> None:
+        self._tasks[task_id] = MemoryTaskInfo(
+            metadata=task_metadata,
+            progress=ProgressReport(actual_value=0.0),
+        )
+
+    async def exists_task(self, task_id: TaskID) -> bool:
+        return task_id in self._tasks
+
+    async def get_task_metadata(self, task_id: TaskID) -> TaskMetadata | None:
+        task_info = self._tasks.get(task_id)
+        if task_info is None:
+            return None
+        return task_info.metadata
+
+    async def get_task_progress(self, task_id: TaskID) -> ProgressReport | None:
+        task_info = self._tasks.get(task_id)
+        if task_info is None:
+            return None
+        return task_info.progress
+
+    async def list_tasks(self, task_context: TaskContext) -> list[Task]:
+        tasks = []
+        task_id_prefix = build_task_id_prefix(task_context)
+        for task_id, task_info in self._tasks.items():
+            if task_id.startswith(task_id_prefix):
+                tasks.append(
+                    Task(
+                        uuid=TaskUUID(task_id[len(task_id_prefix) + 1 :]),
+                        metadata=task_info.metadata,
+                    )
+                )
+        return tasks
+
+    async def remove_task(self, task_id: TaskID) -> None:
+        self._tasks.pop(task_id, None)
+
+    async def set_task_progress(self, task_id: TaskID, report: ProgressReport) -> None:
+        task_info = self._tasks.get(task_id)
+        if task_info is not None:
+            task_info.progress = report

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

@@ -4,7 +4,7 @@
 
 from celery import Celery  # type: ignore[import-untyped]
 from celery.worker.worker import WorkController  # type: ignore[import-untyped]
-from servicelib.base_app_server import BaseAppServer
+from servicelib.base_app_server import STARTUP_TIMEOUT, BaseAppServer
 from servicelib.logging_utils import log_context
 from settings_library.celery import CelerySettings
 
@@ -30,6 +30,8 @@ def _init(startup_complete_event: threading.Event) -> None:
         loop = asyncio.new_event_loop()
         asyncio.set_event_loop(loop)
 
+        shutdown_event = asyncio.Event()
+
         app_server.event_loop = loop
 
         async def _setup():
@@ -46,7 +48,9 @@ async def _setup():
             )
 
         loop.run_until_complete(_setup())
-        loop.run_until_complete(app_server.startup(startup_complete_event))
+        loop.run_until_complete(
+            app_server.startup(startup_complete_event, shutdown_event)
+        )
 
     thread = threading.Thread(
         group=None,
@@ -57,7 +61,7 @@ async def _setup():
     )
     thread.start()
 
-    startup_complete_event.wait()
+    startup_complete_event.wait(STARTUP_TIMEOUT * 1.1)
 
 
 def on_worker_shutdown(sender, **_kwargs) -> None:

packages/celery-library/tests/conftest.py

@@ -0,0 +1,130 @@
+import asyncio
+import datetime
+from collections.abc import AsyncIterator, Callable
+from functools import partial
+from threading import Event
+from typing import Any
+
+import pytest
+from celery import Celery  # type: ignore[import-untyped]
+from celery.contrib.testing.worker import TestWorkController, start_worker
+from celery.signals import worker_init, worker_shutdown
+from celery.worker.worker import WorkController
+from celery_library.backends._memory import MemoryTaskInfoStore
+from celery_library.signals import on_worker_init, on_worker_shutdown
+from celery_library.utils import CeleryTaskManager, get_task_manager
+from pytest_mock import MockerFixture
+from pytest_simcore.helpers.monkeypatch_envs import setenvs_from_dict
+from pytest_simcore.helpers.typing_env import EnvVarsDict
+from servicelib.base_app_server import BaseAppServer
+from settings_library.celery import CelerySettings
+
+pytest_plugins = [
+    "pytest_simcore.environment_configs",
+    "pytest_simcore.repository_paths",
+]
+
+
+class FakeAppServer(BaseAppServer):
+    def __init__(self):
+        self._shutdown_event: asyncio.Event | None = None
+
+    async def startup(
+        self, completed_event: Event, shutdown_event: asyncio.Event
+    ) -> None:
+        completed_event.set()
+        await shutdown_event.wait()
+
+    async def shutdown(self) -> None:
+        if self._shutdown_event is not None:
+            self._shutdown_event.set()
+
+
+@pytest.fixture
+def celery_config() -> dict[str, Any]:
+    return {
+        "broker_connection_retry_on_startup": True,
+        "broker_url": "memory://localhost//",
+        "result_backend": "cache+memory://localhost//",
+        "result_expires": datetime.timedelta(days=7),
+        "result_extended": True,
+        "pool": "threads",
+        "task_default_queue": "default",
+        "task_send_sent_event": True,
+        "task_track_started": True,
+        "worker_send_task_events": True,
+    }
+
+
+@pytest.fixture
+def register_celery_tasks() -> Callable[[Celery], None]:
+    """override if tasks are needed"""
+
+    def _(celery_app: Celery) -> None: ...
+
+    return _
+
+
+@pytest.fixture
+def app_environment(
+    monkeypatch: pytest.MonkeyPatch,
+    env_devel_dict: EnvVarsDict,
+) -> EnvVarsDict:
+    return setenvs_from_dict(
+        monkeypatch,
+        {
+            **env_devel_dict,
+        },
+    )
+
+
+@pytest.fixture
+def celery_settings(
+    app_environment: EnvVarsDict,
+) -> CelerySettings:
+    return CelerySettings.create_from_envs()
+
+
+@pytest.fixture
+async def with_storage_celery_worker(
+    celery_app: Celery,
+    celery_settings: CelerySettings,
+    register_celery_tasks: Callable[[Celery], None],
+    mocker: MockerFixture,
+) -> AsyncIterator[TestWorkController]:
+    mocker.patch(
+        "celery_library.signals.create_task_manager",
+        return_value=CeleryTaskManager(
+            celery_app, celery_settings, MemoryTaskInfoStore()
+        ),
+    )
+
+    def _on_worker_init_wrapper(sender: WorkController, **_kwargs):
+        return partial(on_worker_init, FakeAppServer(), celery_settings)(
+            sender, **_kwargs
+        )
+
+    worker_init.connect(_on_worker_init_wrapper)
+    worker_shutdown.connect(on_worker_shutdown)
+
+    register_celery_tasks(celery_app)
+
+    with start_worker(
+        celery_app,
+        pool="threads",
+        concurrency=1,
+        loglevel="info",
+        perform_ping_check=False,
+        queues="default",
+    ) as worker:
+        yield worker
+
+
+@pytest.fixture
+def celery_task_manager(
+    with_storage_celery_worker: TestWorkController,
+) -> CeleryTaskManager:
+    assert with_storage_celery_worker.app  # nosec
+    assert isinstance(with_storage_celery_worker.app, Celery)  # nosec
+
+    return get_task_manager(with_storage_celery_worker.app)

services/storage/tests/unit/test_modules_celery.py renamed to packages/celery-library/tests/unit/test_tasks.py

@@ -13,7 +13,6 @@
 import pytest
 from celery import Celery, Task
 from celery.contrib.abortable import AbortableTask
-from celery.contrib.testing.worker import TestWorkController
 from celery_library.errors import TransferrableCeleryError
 from celery_library.models import (
     TaskContext,
@@ -28,25 +27,12 @@
 from celery_library.task_manager import CeleryTaskManager
 from celery_library.utils import get_app_server, get_task_manager
 from common_library.errors_classes import OsparcErrorMixin
-from fastapi import FastAPI
 from models_library.progress_bar import ProgressReport
 from servicelib.logging_utils import log_context
-from simcore_service_storage.modules.celery import get_task_manager_from_app
 from tenacity import Retrying, retry_if_exception_type, stop_after_delay, wait_fixed
 
 _logger = logging.getLogger(__name__)
 
-pytest_simcore_core_services_selection = ["postgres", "rabbit"]
-pytest_simcore_ops_services_selection = []
-
-
-@pytest.fixture
-def celery_client(
-    initialized_app: FastAPI,
-    with_storage_celery_worker: TestWorkController,
-) -> CeleryTaskManager:
-    return get_task_manager_from_app(initialized_app)
-
 
 async def _fake_file_processor(
     celery_app: Celery, task_name: str, task_id: str, files: list[str]
@@ -110,11 +96,11 @@ def _(celery_app: Celery) -> None:
 
 
 async def test_submitting_task_calling_async_function_results_with_success_state(
-    celery_client: CeleryTaskManager,
+    celery_task_manager: CeleryTaskManager,
 ):
     task_context = TaskContext(user_id=42)
 
-    task_uuid = await celery_client.submit_task(
+    task_uuid = await celery_task_manager.submit_task(
         TaskMetadata(
             name=fake_file_processor.__name__,
         ),
@@ -128,23 +114,23 @@ async def test_submitting_task_calling_async_function_results_with_success_state
         stop=stop_after_delay(30),
     ):
         with attempt:
-            status = await celery_client.get_task_status(task_context, task_uuid)
+            status = await celery_task_manager.get_task_status(task_context, task_uuid)
             assert status.task_state == TaskState.SUCCESS
 
     assert (
-        await celery_client.get_task_status(task_context, task_uuid)
+        await celery_task_manager.get_task_status(task_context, task_uuid)
     ).task_state == TaskState.SUCCESS
     assert (
-        await celery_client.get_task_result(task_context, task_uuid)
+        await celery_task_manager.get_task_result(task_context, task_uuid)
     ) == "archive.zip"
 
 
 async def test_submitting_task_with_failure_results_with_error(
-    celery_client: CeleryTaskManager,
+    celery_task_manager: CeleryTaskManager,
 ):
     task_context = TaskContext(user_id=42)
 
-    task_uuid = await celery_client.submit_task(
+    task_uuid = await celery_task_manager.submit_task(
         TaskMetadata(
             name=failure_task.__name__,
         ),
@@ -158,49 +144,53 @@ async def test_submitting_task_with_failure_results_with_error(
     ):
 
         with attempt:
-            raw_result = await celery_client.get_task_result(task_context, task_uuid)
+            raw_result = await celery_task_manager.get_task_result(
+                task_context, task_uuid
+            )
             assert isinstance(raw_result, TransferrableCeleryError)
 
-    raw_result = await celery_client.get_task_result(task_context, task_uuid)
+    raw_result = await celery_task_manager.get_task_result(task_context, task_uuid)
     assert f"{raw_result}" == "Something strange happened: BOOM!"
 
 
 async def test_cancelling_a_running_task_aborts_and_deletes(
-    celery_client: CeleryTaskManager,
+    celery_task_manager: CeleryTaskManager,
 ):
     task_context = TaskContext(user_id=42)
 
-    task_uuid = await celery_client.submit_task(
+    task_uuid = await celery_task_manager.submit_task(
         TaskMetadata(
             name=dreamer_task.__name__,
         ),
         task_context=task_context,
     )
 
-    await celery_client.cancel_task(task_context, task_uuid)
+    await celery_task_manager.cancel_task(task_context, task_uuid)
 
     for attempt in Retrying(
         retry=retry_if_exception_type(AssertionError),
         wait=wait_fixed(1),
         stop=stop_after_delay(30),
     ):
         with attempt:
-            progress = await celery_client.get_task_status(task_context, task_uuid)
+            progress = await celery_task_manager.get_task_status(
+                task_context, task_uuid
+            )
             assert progress.task_state == TaskState.ABORTED
 
     assert (
-        await celery_client.get_task_status(task_context, task_uuid)
+        await celery_task_manager.get_task_status(task_context, task_uuid)
     ).task_state == TaskState.ABORTED
 
-    assert task_uuid not in await celery_client.list_tasks(task_context)
+    assert task_uuid not in await celery_task_manager.list_tasks(task_context)
 
 
 async def test_listing_task_uuids_contains_submitted_task(
-    celery_client: CeleryTaskManager,
+    celery_task_manager: CeleryTaskManager,
 ):
     task_context = TaskContext(user_id=42)
 
-    task_uuid = await celery_client.submit_task(
+    task_uuid = await celery_task_manager.submit_task(
         TaskMetadata(
             name=dreamer_task.__name__,
         ),
@@ -213,10 +203,10 @@ async def test_listing_task_uuids_contains_submitted_task(
         stop=stop_after_delay(10),
     ):
         with attempt:
-            tasks = await celery_client.list_tasks(task_context)
+            tasks = await celery_task_manager.list_tasks(task_context)
             assert len(tasks) == 1
             assert task_uuid == tasks[0].uuid
 
-        tasks = await celery_client.list_tasks(task_context)
+        tasks = await celery_task_manager.list_tasks(task_context)
         assert len(tasks) == 1
         assert task_uuid == tasks[0].uuid