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 datetime import datetime, timezone
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)

    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)
        clip_repo = StubClipRepo([])
        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 == "之前的错误"