use hooks

Author: giancarloromeo

packages/celery-library/requirements/_test.in

@@ -11,7 +11,6 @@
 # testing
 coverage
 faker
-fastapi
 httpx
 pint
 pytest

packages/celery-library/tests/conftest.py

@@ -1,11 +1,9 @@
 # pylint: disable=redefined-outer-name
 # pylint: disable=unused-argument
 
-import asyncio
 import datetime
 from collections.abc import AsyncIterator, Callable
 from functools import partial
-from threading import Event
 from typing import Any
 
 import pytest
@@ -32,19 +30,11 @@
 
 
 class FakeAppServer(BaseAppServer):
-    def __init__(self):
-        self._shutdown_event: asyncio.Event | None = None
+    async def on_startup(self) -> None:
+        pass
 
-    async def startup(
-        self, completed_event: Event, shutdown_event: asyncio.Event
-    ) -> None:
-        self._shutdown_event = shutdown_event
-        completed_event.set()
-        await self._shutdown_event.wait()
-
-    async def shutdown(self) -> None:
-        if self._shutdown_event is not None:
-            self._shutdown_event.set()
+    async def on_shutdown(self) -> None:
+        pass
 
 
 @pytest.fixture
@@ -80,13 +70,17 @@ def celery_settings(
     return CelerySettings.create_from_envs()
 
 
+@pytest.fixture
+def app_server() -> BaseAppServer:
+    return FakeAppServer()
+
+
 @pytest.fixture(scope="session")
 def celery_config() -> dict[str, Any]:
     return {
         "broker_connection_retry_on_startup": True,
         "broker_url": "memory://localhost//",
         "result_backend": "cache+memory://localhost//",
-        # "result_backend": celery_settings.CELERY_REDIS_RESULT_BACKEND.build_redis_dsn(RedisDatabase.CELERY_TASKS),
         "result_expires": datetime.timedelta(days=7),
         "result_extended": True,
         "pool": "threads",
@@ -100,13 +94,12 @@ def celery_config() -> dict[str, Any]:
 @pytest.fixture
 async def with_storage_celery_worker(
     celery_app: Celery,
+    app_server: BaseAppServer,
     celery_settings: CelerySettings,
     register_celery_tasks: Callable[[Celery], None],
 ) -> AsyncIterator[TestWorkController]:
     def _on_worker_init_wrapper(sender: WorkController, **_kwargs):
-        return partial(on_worker_init, FakeAppServer(), celery_settings)(
-            sender, **_kwargs
-        )
+        return partial(on_worker_init, app_server, celery_settings)(sender, **_kwargs)
 
     worker_init.connect(_on_worker_init_wrapper)
     worker_shutdown.connect(on_worker_shutdown)

packages/celery-library/tests/unit/test_errors.py

@@ -13,7 +13,7 @@
         AccessRightError(user_id=1, file_id="a/path/to/a/file.txt", location_id=0),
     ],
 )
-def test_workflow(original_error: Exception):
+def test_error(original_error: Exception):
     try:
         raise original_error  # noqa: TRY301
     except Exception as e:  # pylint: disable=broad-exception-caught

packages/service-library/src/servicelib/celery/app_server.py

@@ -17,6 +17,9 @@
 
 
 class BaseAppServer(ABC):
+    def __init__(self) -> None:
+        self._shutdown_event: asyncio.Event | None = None
+
     @property
     def fastapi_app(self) -> "FastAPI":
         raise NotImplementedError
@@ -34,11 +37,23 @@ def event_loop(self, loop: AbstractEventLoop) -> None:
         self._event_loop = loop
 
     @abstractmethod
+    async def on_startup(self) -> None:
+        raise NotImplementedError
+
     async def startup(
         self, completed_event: threading.Event, shutdown_event: asyncio.Event
     ) -> None:
-        pass
+        self._shutdown_event = shutdown_event
+        completed_event.set()
+        await self.on_startup()
+        await self._shutdown_event.wait()
 
     @abstractmethod
-    async def shutdown(self):
-        pass
+    async def on_shutdown(self) -> None:
+        raise NotImplementedError
+
+    async def shutdown(self) -> None:
+        if self._shutdown_event is not None:
+            self._shutdown_event.set()
+
+        await self.on_shutdown()

packages/service-library/src/servicelib/fastapi/celery/app_server.py

@@ -1,5 +1,3 @@
-import asyncio
-import threading
 from datetime import timedelta
 from typing import Final
 
@@ -14,32 +12,24 @@
 
 class FastAPIAppServer(BaseAppServer):
     def __init__(self, app: FastAPI):
+        super().__init__()
         self._app = app
         self._lifespan_manager: LifespanManager | None = None
-        self._shutdown_event: asyncio.Event | None = None
 
     @property
     def fastapi_app(self) -> FastAPI:
         assert isinstance(self._app, FastAPI)  # nosec
         return self._app
 
-    async def startup(
-        self, completed_event: threading.Event, shutdown_event: asyncio.Event
-    ):
+    async def on_startup(self) -> None:
         self._lifespan_manager = LifespanManager(
             self.fastapi_app,
             startup_timeout=_STARTUP_TIMEOUT,
             shutdown_timeout=_SHUTDOWN_TIMEOUT,
         )
-        self._shutdown_event = shutdown_event
         await self._lifespan_manager.__aenter__()
-        completed_event.set()
-        await self._shutdown_event.wait()
-
-    async def shutdown(self):
-        if self._shutdown_event is not None:
-            self._shutdown_event.set()
 
+    async def on_shutdown(self) -> None:
         if self._lifespan_manager is None:
             return
         await self._lifespan_manager.__aexit__(None, None, None)