xiaoxia-saas/tests/unit/test_edit_plan_worker_failure.py

"""P0-2: Celery 任务 render_edit_plan 失败时更新 GenerationTask 状态。

验证：
  - 异常发生时 GenerationTask 状态更新为 failed
  - error_message 记录了异常类型和描述
  - completed_at 被设置
  - 即使 generation_task_id 为空也不崩溃
  - 即使更新 GenerationTask 本身失败也不影响 retry
"""

from __future__ import annotations

import os
import sys
from dataclasses import dataclass, field
from types import ModuleType
from typing import Any, Optional
from unittest.mock import MagicMock, patch

os.environ.setdefault("JWT_SECRET_KEY", "unit-test-secret-key-for-testing")
os.environ.setdefault("DATABASE_URL", "sqlite:///test.db")

import pytest

# ── Mock worker 模块以避免数据库连接 ──────────────────────────────────────────
# worker_app.db 在 import 时会尝试连接数据库，必须在导入 task 模块前 mock

sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "..", "apps", "worker"))
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", ".."))

# 预注册 mock 模块，阻止真实数据库初始化
_mock_db_mod = ModuleType("worker_app.db")
_mock_db_mod.SessionLocal = MagicMock()
sys.modules.setdefault("worker_app.db", _mock_db_mod)

_mock_celery_mod = ModuleType("worker_app.celery_app")
_mock_celery_app = MagicMock()
# 让 @celery_app.task(...) 装饰器透传原始函数，否则函数变成 MagicMock
_mock_celery_app.task = lambda **kwargs: lambda fn: fn
_mock_celery_mod.celery_app = _mock_celery_app
sys.modules.setdefault("worker_app.celery_app", _mock_celery_mod)


# ── Stub domain objects ───────────────────────────────────────────────────────


@dataclass
class _StubStatus:
    value: str

    def __eq__(self, other):
        if isinstance(other, str):
            return self.value == other
        if isinstance(other, _StubStatus):
            return self.value == other.value
        return NotImplemented


@dataclass
class StubEditPlan:
    id: str = "plan-001"
    template_id: str = "tmpl-001"
    status: Any = None
    config: dict = field(default_factory=dict)
    project_id: str = ""
    created_by_user_id: str = "user-001"

    def mark_failed(self):
        self.status = _StubStatus("failed")

    def mark_completed(self):
        self.status = _StubStatus("completed")


@dataclass
class StubGenerationTask:
    id: str = "gen-task-001"
    status: Any = field(default_factory=lambda: _StubStatus("pending"))
    error_message: str = ""
    progress: float = 0.0
    result_count: int = 0
    started_at: Any = None
    completed_at: Any = None
    project_id: str = ""
    created_by_user_id: str = "user-001"


@dataclass
class StubClip:
    id: str = "clip-001"
    plan_id: str = "plan-001"
    asset_id: str = "assets/video.mp4"
    order: int = 1
    status: Any = field(default_factory=lambda: _StubStatus("ready"))
    transition_effect: str = ""
    text_content: str = ""
    clip_type: str = "MAIN"
    duration: float = 0.0

    def mark_failed(self):
        self.status = _StubStatus("failed")

    def mark_rendered(self):
        self.status = _StubStatus("rendered")


# ── Stub repositories ─────────────────────────────────────────────────────────


class StubPlanRepo:
    def __init__(self, plan: StubEditPlan):
        self._plan = plan

    def get(self, plan_id: str) -> Optional[StubEditPlan]:
        if plan_id == self._plan.id:
            return self._plan
        return None

    def update(self, plan: StubEditPlan) -> StubEditPlan:
        self._plan = plan
        return plan


class StubClipRepo:
    def __init__(self, clips: list[StubClip] | None = None):
        self._clips = clips or []

    def list_by_plan(self, plan_id: str, skip: int = 0, limit: int = 10000) -> list[StubClip]:
        return [c for c in self._clips if c.plan_id == plan_id]

    def get(self, clip_id: str) -> Optional[StubClip]:
        for c in self._clips:
            if c.id == clip_id:
                return c
        return None

    def update(self, clip: StubClip) -> StubClip:
        return clip


class StubGenTaskRepo:
    def __init__(self, task: StubGenerationTask | None = None):
        self._store: dict[str, StubGenerationTask] = {}
        if task:
            self._store[task.id] = task

    def get(self, task_id: str) -> Optional[StubGenerationTask]:
        return self._store.get(task_id)

    def update(self, task: StubGenerationTask) -> StubGenerationTask:
        self._store[task.id] = task
        return task


# ── Import task module (after mocks are in place) ─────────────────────────────

from worker_app.tasks.edit_plan_generation import render_edit_plan

# ── Tests ─────────────────────────────────────────────────────────────────────


class TestRenderEditPlanFailureUpdatesGenTask:
    """P0-2: render_edit_plan 异常时更新 GenerationTask 状态为 failed"""

    def _make_bound_task(self):
        """构建绑定的 Celery task mock"""
        task = MagicMock()
        task.retry = MagicMock(side_effect=RuntimeError("retry called"))
        return task

    def test_exception_marks_gen_task_failed(self):
        """异常时 GenerationTask.status 被设为 failed"""
        plan = StubEditPlan(status=_StubStatus("rendering"))
        plan.config["generation_task_id"] = "gen-task-001"
        gen_task = StubGenerationTask(id="gen-task-001", status=_StubStatus("running"))

        plan_repo = StubPlanRepo(plan)
        gen_task_repo = StubGenTaskRepo(gen_task)

        # 让 clip_repo 抛异常以触发 except 路径
        clip_repo_bad = MagicMock()
        clip_repo_bad.list_by_plan.side_effect = RuntimeError("OSS 连接失败")

        def fake_get_repos():
            yield plan_repo, clip_repo_bad, gen_task_repo, MagicMock()

        bound_task = self._make_bound_task()

        with patch(
            "worker_app.tasks.edit_plan_generation._get_repos",
            side_effect=fake_get_repos,
        ):
            with pytest.raises(RuntimeError, match="retry called"):
                render_edit_plan(bound_task, "plan-001")

        # 核心断言：GenerationTask 状态为 failed（生产代码赋值为字符串）
        assert gen_task.status == "failed"

    def test_exception_records_error_message(self):
        """异常时 error_message 包含异常类型和描述"""
        plan = StubEditPlan(status=_StubStatus("rendering"))
        plan.config["generation_task_id"] = "gen-task-001"
        gen_task = StubGenerationTask(id="gen-task-001", status=_StubStatus("running"))

        plan_repo = StubPlanRepo(plan)
        clip_repo = MagicMock()
        clip_repo.list_by_plan.side_effect = RuntimeError("DB 查询超时")
        gen_task_repo = StubGenTaskRepo(gen_task)

        def fake_get_repos():
            yield plan_repo, clip_repo, gen_task_repo, MagicMock()

        bound_task = self._make_bound_task()

        with patch(
            "worker_app.tasks.edit_plan_generation._get_repos",
            side_effect=fake_get_repos,
        ):
            with pytest.raises(RuntimeError, match="retry called"):
                render_edit_plan(bound_task, "plan-001")

        assert gen_task.status == "failed"
        assert "DB 查询超时" in gen_task.error_message
        assert "RuntimeError" in gen_task.error_message

    def test_exception_sets_completed_at(self):
        """异常时 completed_at 被设置"""
        plan = StubEditPlan(status=_StubStatus("rendering"))
        plan.config["generation_task_id"] = "gen-task-001"
        gen_task = StubGenerationTask(id="gen-task-001", status=_StubStatus("running"))

        plan_repo = StubPlanRepo(plan)
        clip_repo = MagicMock()
        clip_repo.list_by_plan.side_effect = RuntimeError("boom")
        gen_task_repo = StubGenTaskRepo(gen_task)

        def fake_get_repos():
            yield plan_repo, clip_repo, gen_task_repo, MagicMock()

        bound_task = self._make_bound_task()

        with patch(
            "worker_app.tasks.edit_plan_generation._get_repos",
            side_effect=fake_get_repos,
        ):
            with pytest.raises(RuntimeError):
                render_edit_plan(bound_task, "plan-001")

        assert gen_task.completed_at is not None

    def test_no_generation_task_id_does_not_crash(self):
        """generation_task_id 为空时，异常处理不崩溃"""
        plan = StubEditPlan(status=_StubStatus("rendering"))
        plan.config = {}  # 不设置 generation_task_id

        plan_repo = StubPlanRepo(plan)
        clip_repo = MagicMock()
        clip_repo.list_by_plan.side_effect = RuntimeError("boom")
        gen_task_repo = StubGenTaskRepo()  # 空 repo

        def fake_get_repos():
            yield plan_repo, clip_repo, gen_task_repo, MagicMock()

        bound_task = self._make_bound_task()

        with patch(
            "worker_app.tasks.edit_plan_generation._get_repos",
            side_effect=fake_get_repos,
        ):
            with pytest.raises(RuntimeError, match="retry called"):
                render_edit_plan(bound_task, "plan-001")

        # 计划仍被标记为 failed
        assert plan.status.value == "failed"

    def test_gen_task_update_failure_does_not_block_retry(self):
        """更新 GenerationTask 失败时，不影响 retry 流程"""
        plan = StubEditPlan(status=_StubStatus("rendering"))
        plan.config["generation_task_id"] = "gen-task-001"
        gen_task = StubGenerationTask(id="gen-task-001", status=_StubStatus("running"))

        plan_repo = StubPlanRepo(plan)
        clip_repo = MagicMock()
        clip_repo.list_by_plan.side_effect = RuntimeError("原始错误")
        # gen_task_repo.update 也抛异常
        gen_task_repo = MagicMock()
        gen_task_repo.get.return_value = gen_task
        gen_task_repo.update.side_effect = RuntimeError("DB 写入失败")

        def fake_get_repos():
            yield plan_repo, clip_repo, gen_task_repo, MagicMock()

        bound_task = self._make_bound_task()

        with patch(
            "worker_app.tasks.edit_plan_generation._get_repos",
            side_effect=fake_get_repos,
        ):
            with pytest.raises(RuntimeError, match="retry called"):
                render_edit_plan(bound_task, "plan-001")

        # retry 被调用说明流程正确
        bound_task.retry.assert_called_once()

    def test_already_failed_gen_task_not_overwritten(self):
        """已经 failed 的 GenerationTask 不会被重复更新"""
        plan = StubEditPlan(status=_StubStatus("rendering"))
        plan.config["generation_task_id"] = "gen-task-001"
        gen_task = StubGenerationTask(
            id="gen-task-001",
            status=_StubStatus("failed"),  # 已经是 failed
            error_message="之前的错误",
        )

        plan_repo = StubPlanRepo(plan)
        clip_repo = MagicMock()
        clip_repo.list_by_plan.side_effect = RuntimeError("新错误")
        gen_task_repo = StubGenTaskRepo(gen_task)

        def fake_get_repos():
            yield plan_repo, clip_repo, gen_task_repo, MagicMock()

        bound_task = self._make_bound_task()

        with patch(
            "worker_app.tasks.edit_plan_generation._get_repos",
            side_effect=fake_get_repos,
        ):
            with pytest.raises(RuntimeError, match="retry called"):
                render_edit_plan(bound_task, "plan-001")

        # error_message 应保持原值，不被覆盖
        assert gen_task.error_message == "之前的错误"