from __future__ import annotations import asyncio import inspect from types import SimpleNamespace import pytest from sqlalchemy import create_engine from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine from sqlalchemy.orm import Session, sessionmaker from app.core.db import Base from app.core.task_manager import DeliveryMode, SqlAlchemyTaskStore, TaskManager from app.core.task_manager import SyncSqlAlchemyTaskStore from app.core.task_manager.types import TaskStatus from app.models.task import GenerationTask from app.tasks import execute_task as execute_task_module from app.services.worker.task_executor import AbstractAsyncDelegatingExecutor, AbstractWorkerTaskExecutor, WorkerTaskContext from app.services.worker.task_registry import task_executor_registry class _DummyTask: async def run(self, *args, **kwargs): # noqa: ANN002, ANN003 return None async def status(self) -> dict: return {} async def is_done(self) -> bool: return False async def get_result(self): return None @pytest.mark.asyncio async def test_task_manager_create_persists_task_kind() -> None: engine = create_async_engine("sqlite+aiosqlite:///:memory:", future=True) session_local = async_sessionmaker(engine, class_=AsyncSession, expire_on_commit=False) import app.models.task # noqa: F401 async with engine.begin() as conn: await conn.run_sync(Base.metadata.create_all) async with session_local() as db: store = SqlAlchemyTaskStore(db) tm = TaskManager(store=store, strategies={}) record = await tm.create( task=_DummyTask(), mode=DeliveryMode.async_polling, task_kind="script_divide", run_args={"chapter_id": "chapter-1"}, ) row = await db.get(GenerationTask, record.id) assert row is not None assert row.task_kind == "script_divide" assert row.payload["task_kind"] == "script_divide" await engine.dispose() def test_run_task_celery_routes_by_task_kind(monkeypatch, tmp_path) -> None: db_path = tmp_path / "task-execute.db" sync_engine = create_engine(f"sqlite:///{db_path}", future=True) sync_session_local = sessionmaker(sync_engine, class_=Session, expire_on_commit=False) import app.models.task # noqa: F401 Base.metadata.create_all(sync_engine) with sync_session_local() as db: db.add( GenerationTask( id="task-1", mode="async_polling", task_kind="script_divide", status="pending", progress=0, payload={"task_kind": "script_divide", "run_args": {}}, result=None, error="", ) ) db.commit() called: list[str] = [] class _FakeExecutor: def run(self, task_id: str) -> None: called.append(task_id) def _fake_resolve(task_kind: str): assert task_kind == "script_divide" return _FakeExecutor() monkeypatch.setattr(execute_task_module, "sync_session_maker", sync_session_local) monkeypatch.setattr(execute_task_module.task_executor_registry, "resolve", _fake_resolve) execute_task_module.run_task_celery("task-1") assert called == ["task-1"] sync_engine.dispose() def test_enqueue_task_execution_records_executor(monkeypatch, tmp_path) -> None: db_path = tmp_path / "task-enqueue.db" sync_engine = create_engine(f"sqlite:///{db_path}", future=True) sync_session_local = sessionmaker(sync_engine, class_=Session, expire_on_commit=False) import app.models.task # noqa: F401 Base.metadata.create_all(sync_engine) with sync_session_local() as db: db.add( GenerationTask( id="task-2", mode="async_polling", task_kind="script_extract", status="pending", progress=0, payload={"task_kind": "script_extract", "run_args": {}}, result=None, error="", ) ) db.commit() monkeypatch.setattr(execute_task_module, "sync_session_maker", sync_session_local) monkeypatch.setattr( execute_task_module.run_task_celery, "delay", lambda task_id: SimpleNamespace(id=f"celery-{task_id}"), ) result = execute_task_module.enqueue_task_execution("task-2") assert result.id == "celery-task-2" with sync_session_local() as db: row = db.get(GenerationTask, "task-2") assert row is not None assert row.executor_type == "celery" assert row.executor_task_id == "celery-task-2" sync_engine.dispose() def test_revoke_task_execution_revokes_celery_task(monkeypatch, tmp_path) -> None: db_path = tmp_path / "task-revoke.db" sync_engine = create_engine(f"sqlite:///{db_path}", future=True) sync_session_local = sessionmaker(sync_engine, class_=Session, expire_on_commit=False) import app.models.task # noqa: F401 Base.metadata.create_all(sync_engine) with sync_session_local() as db: db.add( GenerationTask( id="task-revoke", mode="async_polling", task_kind="video_generation", status="running", progress=30, payload={"task_kind": "video_generation", "run_args": {}}, result=None, error="", executor_type="celery", executor_task_id="celery-task-revoke", ) ) db.commit() calls: list[tuple[str, bool, str]] = [] class _FakeAsyncResult: def __init__(self, task_id: str, app=None) -> None: # noqa: ANN001 self.task_id = task_id def revoke(self, *, terminate: bool, signal: str) -> None: calls.append((self.task_id, terminate, signal)) monkeypatch.setattr(execute_task_module, "sync_session_maker", sync_session_local) monkeypatch.setattr(execute_task_module, "AsyncResult", _FakeAsyncResult) assert execute_task_module.revoke_task_execution("task-revoke") is True assert calls == [("celery-task-revoke", True, "SIGTERM")] sync_engine.dispose() def test_async_delegating_executors_use_positional_runner_signature() -> None: for task_kind, executor in task_executor_registry._executors.items(): if not isinstance(executor, AbstractAsyncDelegatingExecutor): continue signature = inspect.signature(executor._runner) params = list(signature.parameters.values()) assert len(params) == 2, f"{task_kind} runner must accept exactly 2 params, got {signature}" assert all( param.kind in (inspect.Parameter.POSITIONAL_ONLY, inspect.Parameter.POSITIONAL_OR_KEYWORD) for param in params ), f"{task_kind} runner must accept positional args, got {signature}" def test_async_delegating_executor_marks_failed_on_timeout(monkeypatch, tmp_path) -> None: db_path = tmp_path / "async-executor-timeout.db" sync_engine = create_engine(f"sqlite:///{db_path}", future=True) sync_session_local = sessionmaker(sync_engine, class_=Session, expire_on_commit=False) import app.models.task # noqa: F401 Base.metadata.create_all(sync_engine) with sync_session_local() as db: db.add( GenerationTask( id="task-timeout", mode="async_polling", task_kind="video_generation", status="pending", progress=0, payload={"task_kind": "video_generation", "run_args": {}}, result=None, error="", ) ) db.commit() async def _runner(_task_id: str, _run_args: dict) -> None: await asyncio.sleep(0.05) async def _close_db() -> None: return None monkeypatch.setattr("app.services.worker.task_executor.reset_db_runtime", lambda: None) monkeypatch.setattr("app.services.worker.task_executor.close_db", _close_db) executor = AbstractAsyncDelegatingExecutor( task_kind="video_generation", runner=_runner, timeout_seconds=0.01, session_maker=sync_session_local, ) executor.run("task-timeout") with sync_session_local() as db: row = db.get(GenerationTask, "task-timeout") assert row is not None assert row.status == "failed" assert "timed out" in (row.error or "") sync_engine.dispose() def test_async_delegating_executor_marks_cancelled_before_start(monkeypatch, tmp_path) -> None: db_path = tmp_path / "async-executor-cancel.db" sync_engine = create_engine(f"sqlite:///{db_path}", future=True) sync_session_local = sessionmaker(sync_engine, class_=Session, expire_on_commit=False) import app.models.task # noqa: F401 Base.metadata.create_all(sync_engine) with sync_session_local() as db: db.add( GenerationTask( id="task-cancelled", mode="async_polling", task_kind="image_generation", status="running", progress=10, payload={"task_kind": "image_generation", "run_args": {}}, result=None, error="", cancel_requested=True, ) ) db.commit() called: list[str] = [] async def _runner(task_id: str, _run_args: dict) -> None: called.append(task_id) async def _close_db() -> None: return None monkeypatch.setattr("app.services.worker.task_executor.reset_db_runtime", lambda: None) monkeypatch.setattr("app.services.worker.task_executor.close_db", _close_db) executor = AbstractAsyncDelegatingExecutor( task_kind="image_generation", runner=_runner, timeout_seconds=10, session_maker=sync_session_local, ) executor.run("task-cancelled") assert called == [] with sync_session_local() as db: row = db.get(GenerationTask, "task-cancelled") assert row is not None assert row.status == "cancelled" assert bool(row.cancel_requested) is True sync_engine.dispose() def test_async_delegating_executor_closes_db_in_same_event_loop(monkeypatch, tmp_path) -> None: db_path = tmp_path / "async-executor-close-loop.db" sync_engine = create_engine(f"sqlite:///{db_path}", future=True) sync_session_local = sessionmaker(sync_engine, class_=Session, expire_on_commit=False) import app.models.task # noqa: F401 Base.metadata.create_all(sync_engine) with sync_session_local() as db: db.add( GenerationTask( id="task-close-loop", mode="async_polling", task_kind="image_generation", status="pending", progress=0, payload={"task_kind": "image_generation", "run_args": {}}, result=None, error="", ) ) db.commit() loops: dict[str, asyncio.AbstractEventLoop] = {} async def _runner(_task_id: str, _run_args: dict) -> None: loops["runner"] = asyncio.get_running_loop() async def _close_db() -> None: loops["close_db"] = asyncio.get_running_loop() monkeypatch.setattr("app.services.worker.task_executor.reset_db_runtime", lambda: None) monkeypatch.setattr("app.services.worker.task_executor.close_db", _close_db) executor = AbstractAsyncDelegatingExecutor( task_kind="image_generation", runner=_runner, timeout_seconds=10, session_maker=sync_session_local, ) executor.run("task-close-loop") assert loops["runner"] is loops["close_db"] sync_engine.dispose() def test_sync_executor_marks_failed_on_boundary_timeout(tmp_path) -> None: db_path = tmp_path / "sync-executor-timeout.db" sync_engine = create_engine(f"sqlite:///{db_path}", future=True) sync_session_local = sessionmaker(sync_engine, class_=Session, expire_on_commit=False) import app.models.task # noqa: F401 Base.metadata.create_all(sync_engine) with sync_session_local() as db: db.add( GenerationTask( id="task-sync-timeout", mode="async_polling", task_kind="script_divide", status="pending", progress=0, payload={"task_kind": "script_divide", "run_args": {}}, result=None, error="", ) ) db.commit() class _SlowExecutor(AbstractWorkerTaskExecutor): task_kind = "script_divide" timeout_seconds = 0.01 def __init__(self) -> None: super().__init__(session_maker=sync_session_local) def execute(self, ctx: WorkerTaskContext, run_args: dict[str, object]) -> dict[str, object]: # noqa: ARG002 import time time.sleep(0.03) return {"ok": True} _SlowExecutor().run("task-sync-timeout") with sync_session_local() as db: row = db.get(GenerationTask, "task-sync-timeout") assert row is not None assert row.status == "failed" assert "timed out" in (row.error or "") sync_engine.dispose() def test_task_executor_registry_resolves_sync_and_async_executor_types() -> None: divide_executor = task_executor_registry.resolve("script_divide") video_executor = task_executor_registry.resolve("video_generation") assert isinstance(divide_executor, AbstractWorkerTaskExecutor) assert divide_executor.task_kind == "script_divide" assert isinstance(video_executor, AbstractAsyncDelegatingExecutor) assert video_executor.task_kind == "video_generation"