[DOP-21799] Refactor celery initialization

Author: Ilyas Gasanov

syncmaster/backend/__init__.py

@@ -1,5 +1,6 @@
 # SPDX-FileCopyrightText: 2023-2024 MTS PJSC
 # SPDX-License-Identifier: Apache-2.0
+from celery import Celery
 from fastapi import FastAPI, HTTPException
 from fastapi.exceptions import RequestValidationError
 from pydantic import ValidationError
@@ -20,6 +21,15 @@
 from syncmaster.exceptions import SyncmasterError
 
 
+def celery_factory(settings: Settings) -> Celery:
+    app = Celery(
+        __name__,
+        broker=settings.broker.url,
+        backend="db+" + settings.database.sync_url,
+    )
+    return app
+
+
 def application_factory(settings: Settings) -> FastAPI:
     application = FastAPI(
         title="Syncmaster",
@@ -30,6 +40,7 @@ def application_factory(settings: Settings) -> FastAPI:
         redoc_url=None,
     )
     application.state.settings = settings
+    application.state.celery = celery_factory(settings)
     application.include_router(api_router)
     application.exception_handler(RequestValidationError)(validation_exception_handler)
     application.exception_handler(ValidationError)(validation_exception_handler)
@@ -44,6 +55,7 @@ def application_factory(settings: Settings) -> FastAPI:
         {
             Settings: lambda: settings,
             UnitOfWork: get_uow(session_factory, settings=settings),
+            Celery: lambda: application.state.celery,
         },
     )
 

syncmaster/backend/api/v1/runs.py

@@ -5,11 +5,11 @@
 from typing import Annotated
 
 from asgi_correlation_id import correlation_id
+from celery import Celery
 from fastapi import APIRouter, Depends, Query
 from jinja2 import Template
 from kombu.exceptions import KombuError
 
-from syncmaster.backend.celery import app as celery
 from syncmaster.backend.dependencies import Stub
 from syncmaster.backend.services import UnitOfWork, get_user
 from syncmaster.backend.settings import ServerAppSettings as Settings
@@ -84,6 +84,7 @@ async def read_run(
 async def start_run(
     create_run_data: CreateRunSchema,
     settings: Annotated[Settings, Depends(Stub(Settings))],
+    celery: Annotated[Celery, Depends(Stub(Celery))],
     unit_of_work: UnitOfWork = Depends(UnitOfWork),
     current_user: User = Depends(get_user(is_active=True)),
 ) -> ReadRunSchema:

syncmaster/backend/celery.py

@@ -1,4 +1,14 @@
 # SPDX-FileCopyrightText: 2023-2024 MTS PJSC
 # SPDX-License-Identifier: Apache-2.0
-from syncmaster.scheduler.transfer_fetcher import TransferFetcher
-from syncmaster.scheduler.transfer_job_manager import TransferJobManager
+from celery import Celery
+
+from syncmaster.scheduler.settings import SchedulerAppSettings
+
+
+def celery_factory(settings: SchedulerAppSettings) -> Celery:
+    app = Celery(
+        __name__,
+        broker=settings.broker.url,
+        backend="db+" + settings.database.sync_url,
+    )
+    return app

syncmaster/scheduler/__init__.py

@@ -1,6 +1,7 @@
 # SPDX-FileCopyrightText: 2023-2024 MTS PJSC
 # SPDX-License-Identifier: Apache-2.0
+from syncmaster.scheduler import celery_factory
 from syncmaster.scheduler.settings import SchedulerAppSettings
-from syncmaster.worker import celery_factory
 
+# Global object, since the TransferJobManager.send_job_to_celery method is static
 app = celery_factory(SchedulerAppSettings())

syncmaster/scheduler/celery.py

@@ -50,8 +50,11 @@ def update_jobs(self, transfers: list[Transfer]) -> None:
     @staticmethod
     async def send_job_to_celery(transfer_id: int) -> None:
         """
-        Do not pass additional arguments like settings,
+        1. Do not pass additional arguments like settings,
         otherwise they will be serialized in jobs table.
+        2. Instance methods are bound to specific objects and cannot be reliably serialized
+        due to the weak reference problem. Use a static method instead, as it is not
+        object-specific and can be serialized.
         """
         settings = Settings()
 

syncmaster/scheduler/transfer_job_manager.py

@@ -3,9 +3,10 @@
 from celery import Celery
 
 from syncmaster.worker.base import WorkerTask
+from syncmaster.worker.settings import WorkerAppSettings
 
 
-def celery_factory(settings) -> Celery:
+def celery_factory(settings: WorkerAppSettings) -> Celery:
     app = Celery(
         __name__,
         broker=settings.broker.url,

syncmaster/worker/__init__.py

@@ -4,11 +4,13 @@
 import time
 from collections.abc import AsyncGenerator, Callable
 from pathlib import Path
+from unittest.mock import AsyncMock, Mock
 
 import pytest
 import pytest_asyncio
 from alembic.config import Config as AlembicConfig
 from celery import Celery
+from fastapi import FastAPI
 from httpx import AsyncClient
 from sqlalchemy.ext.asyncio import (
     AsyncEngine,
@@ -129,6 +131,26 @@ async def session(sessionmaker: async_sessionmaker[AsyncSession]):
         await session.close()
 
 
+@pytest.fixture(scope="session")
+def mocked_celery() -> Celery:
+    celery_app = Mock(Celery)
+    celery_app.send_task = AsyncMock()
+    return celery_app
+
+
+@pytest_asyncio.fixture(scope="session")
+async def app(settings: Settings, mocked_celery: Celery) -> FastAPI:
+    app = application_factory(settings=settings)
+    app.dependency_overrides[Celery] = lambda: mocked_celery
+    return app
+
+
+@pytest_asyncio.fixture(scope="session")
+async def client_with_mocked_celery(app: FastAPI) -> AsyncGenerator:
+    async with AsyncClient(app=app, base_url="http://testserver") as client:
+        yield client
+
+
 @pytest_asyncio.fixture(scope="session")
 async def client(settings: Settings) -> AsyncGenerator:
     logger.info("START CLIENT FIXTURE")

tests/conftest.py

@@ -1,3 +1,3 @@
-from syncmaster.scheduler.celery import app as celery
+from syncmaster.worker.celery import app as celery
 
 celery.conf.update(imports=list(celery.conf.imports) + ["tests.test_integration.test_scheduler.test_task"])

tests/test_integration/celery_test.py

@@ -7,7 +7,8 @@
 
 from syncmaster.backend.settings import ServerAppSettings as Settings
 from syncmaster.db.models import Run, Status
-from syncmaster.scheduler import TransferFetcher, TransferJobManager
+from syncmaster.scheduler.transfer_fetcher import TransferFetcher
+from syncmaster.scheduler.transfer_job_manager import TransferJobManager
 from tests.mocks import MockTransfer
 
 pytestmark = [pytest.mark.asyncio, pytest.mark.worker, pytest.mark.scheduler_integration]