import contextlib

from fastapi.applications import AppType

from task_management import broker


@contextlib.asynccontextmanager
async def lifespan(app: AppType):
    if not broker.is_worker_process:
        await broker.startup()

    yield

    if not broker.is_worker_process:
        await broker.shutdown()