xiaoxia-saas/tests/unit/test_task_fault_tolerance_1709.py

"""Issue #1709 任务容错：孤儿任务恢复 + 429 限流结构化提示。

覆盖：
1. 仓储层：count_running_by_user/count_running_total 计数正确（预览/正式任务都计入）
2. 仓储层：estimate_avg_duration_seconds 耗时估算（有历史/无历史）
3. 限流核心：build_rate_limit_detail 返回结构化 code/message/排队数/预计等待
4. worker 侧：cleanup_stale_running/pending 核心函数——中断任务被重置为 failed
   且原因写明（容器重启/超时中断），正常任务不受影响
"""

import sys
from datetime import datetime, timedelta, timezone
from pathlib import Path
from unittest.mock import MagicMock

sys.path.insert(0, str(Path(__file__).resolve().parents[2] / "apps" / "api"))
sys.path.insert(0, str(Path(__file__).resolve().parents[2] / "apps" / "worker"))

# 预注入 mock worker_app.db，防止真实数据库连接初始化（与其他 worker 测试同模式）
_mock_db = MagicMock()
_mock_db.SessionLocal = MagicMock()
sys.modules.setdefault("worker_app.db", _mock_db)

from app.core import task_enqueue  # noqa: E402
from sqlalchemy import create_engine, text  # noqa: E402
from sqlalchemy.orm import sessionmaker  # noqa: E402
from worker_app.tasks import _startup  # noqa: E402

from packages.adapters.sqlalchemy_impl.generation_task_repository import (  # noqa: E402
    SQLAlchemyGenerationTaskRepository,
)
from packages.adapters.sqlalchemy_impl.models import Base  # noqa: E402
from packages.domain import GenerationTask, GenerationTaskStatus  # noqa: E402


def _repository():
    engine = create_engine("sqlite:///:memory:", connect_args={"check_same_thread": False})
    Base.metadata.create_all(engine)
    session = sessionmaker(bind=engine)()
    return SQLAlchemyGenerationTaskRepository(session), session, engine


def _make_task(**kwargs) -> GenerationTask:
    defaults = dict(
        project_id="proj-1",
        asset_library_id="lib-1",
        created_by_user_id="user-1",
    )
    defaults.update(kwargs)
    return GenerationTask.create(**defaults)


def _age_task(engine, task_id, *, updated_minutes=None, created_minutes=None):
    """用 SQL 直接把 updated_at/created_at 改到过去（模拟孤儿任务）。"""
    sets, params = [], {"id": task_id}
    if updated_minutes is not None:
        sets.append("updated_at = :uts")
        params["uts"] = datetime.now(timezone.utc) - timedelta(minutes=updated_minutes)
    if created_minutes is not None:
        sets.append("created_at = :cts")
        params["cts"] = datetime.now(timezone.utc) - timedelta(minutes=created_minutes)
    with engine.connect() as conn:
        conn.execute(text(f"UPDATE generation_tasks SET {', '.join(sets)} WHERE id = :id"), params)
        conn.commit()


# ---------------------------------------------------------------------------
# 1. running 计数（限流"渲染中"数量）
# ---------------------------------------------------------------------------


def test_count_running_by_user_mix_statuses():
    """count_running_by_user 只统计该用户 running，不含 pending/completed/failed。"""
    repo, _, _ = _repository()
    t1 = _make_task(project_id="p1")
    repo.create(t1)  # pending
    t2 = _make_task(project_id="p2")
    repo.create(t2)
    t2.mark_processing()
    repo.update(t2)
    t3 = _make_task(project_id="p3")
    repo.create(t3)
    t3.mark_processing()
    repo.update(t3)
    t4 = _make_task(project_id="p4")
    repo.create(t4)
    t4.mark_processing()
    repo.update(t4)
    t4.mark_completed()
    repo.update(t4)
    t5 = _make_task(project_id="p5", created_by_user_id="user-2")
    repo.create(t5)
    t5.mark_processing()
    repo.update(t5)

    assert repo.count_running_by_user("user-1") == 2
    assert repo.count_running_by_user("user-2") == 1
    assert repo.count_running_total() == 3


def test_count_running_total_empty():
    repo, _, _ = _repository()
    assert repo.count_running_total() == 0
    assert repo.count_running_by_user("nobody") == 0


def test_preview_tasks_counted_in_running():
    """预览任务（is_preview=True，工单实测卡 80% 的那种）同样计入 running。"""
    repo, _, _ = _repository()
    t = _make_task(is_preview=True)
    repo.create(t)
    t.mark_processing()
    repo.update(t)
    assert repo.count_running_by_user("user-1") == 1
    assert repo.count_running_total() == 1


# ---------------------------------------------------------------------------
# 2. 平均耗时估算（429 等待预估依据）
# ---------------------------------------------------------------------------


def _complete_task(repo, engine, task, duration_seconds: float):
    repo.create(task)
    task.mark_processing()
    repo.update(task)
    task.mark_completed()
    repo.update(task)
    now = datetime.now(timezone.utc)
    with engine.connect() as conn:
        conn.execute(
            text("UPDATE generation_tasks SET started_at = :s, completed_at = :c WHERE id = :id"),
            {"s": now - timedelta(seconds=duration_seconds), "c": now, "id": task.id},
        )
        conn.commit()


def test_estimate_avg_duration_with_history():
    """有历史完成任务时返回平均耗时（秒）。"""
    repo, _, engine = _repository()
    _complete_task(repo, engine, _make_task(project_id="p1"), 60.0)
    _complete_task(repo, engine, _make_task(project_id="p2"), 180.0)

    avg = repo.estimate_avg_duration_seconds(default_seconds=120.0)
    assert 119.0 < avg < 121.0  # (60+180)/2 = 120


def test_estimate_avg_duration_no_history_returns_default():
    """无历史数据时返回默认值。"""
    repo, _, _ = _repository()
    assert repo.estimate_avg_duration_seconds(default_seconds=90.0) == 90.0


# ---------------------------------------------------------------------------
# 3. build_rate_limit_detail 结构化提示（前端区分"排队"与"创建失败"）
# ---------------------------------------------------------------------------


def test_user_rate_limit_detail_structure():
    """429 用户限流：返回 USER_QUEUE_FULL + 排队/渲染数 + 预计等待。"""
    repo, _, _ = _repository()
    for i in range(2):  # 2 个渲染中
        t = _make_task(project_id=f"rp{i}")
        repo.create(t)
        t.mark_processing()
        repo.update(t)

    exc = task_enqueue.UserPendingLimitExceeded(user_id="user-1", pending_count=3, limit=3)
    detail = task_enqueue.build_rate_limit_detail(exc, repo, scope="user")

    assert detail["code"] == task_enqueue.ERROR_CODE_USER_QUEUE_FULL
    assert detail["queued_count"] == 3
    assert detail["running_count"] == 2
    assert detail["limit"] == 3
    assert detail["estimated_wait_seconds"] > 0
    assert "排队" in detail["message"]
    assert "user-1" not in detail["message"]  # 不泄露内部 ID


def test_global_rate_limit_detail_structure():
    """503 全局繁忙：返回 SYSTEM_QUEUE_FULL。"""
    repo, _, _ = _repository()
    exc = task_enqueue.GlobalQueueFull(pending_count=20, limit=20)
    detail = task_enqueue.build_rate_limit_detail(exc, repo, scope="global")

    assert detail["code"] == task_enqueue.ERROR_CODE_SYSTEM_QUEUE_FULL
    assert detail["queued_count"] == 20
    assert detail["limit"] == 20
    assert detail["estimated_wait_seconds"] > 0
    assert "系统繁忙" in detail["message"]


def test_wait_estimate_uses_concurrency():
    """等待预估：排队 8 个 / 并发 4 = 2 批 × 平均耗时。"""

    class FakeRepo:
        def estimate_avg_duration_seconds(self, limit=20, default_seconds=120.0):
            return 100.0

    wait = task_enqueue._estimate_wait_seconds(8, FakeRepo())
    assert wait == 200  # ceil(8/4)=2 批 × 100 秒


def test_wait_estimate_repo_without_methods_uses_default():
    """仓储没有新方法（旧 mock/鸭子类型）时用默认 120 秒兜底，不抛错。"""

    class LegacyRepo:
        """只实现旧接口的仓储（模拟未升级的调用方）。"""

        def count_pending_total(self):
            return 0

    wait = task_enqueue._estimate_wait_seconds(4, LegacyRepo())
    assert wait == 120  # ceil(4/4)=1 批 × 120 默认


def test_rate_limit_detail_running_count_falls_back_to_zero():
    """仓储不支持 running 计数时，running_count 优雅降级为 0。"""

    class LegacyRepo:
        def count_pending_total(self):
            return 0

    exc = task_enqueue.GlobalQueueFull(pending_count=20, limit=20)
    detail = task_enqueue.build_rate_limit_detail(exc, LegacyRepo(), scope="global")
    assert detail["running_count"] == 0
    assert detail["code"] == task_enqueue.ERROR_CODE_SYSTEM_QUEUE_FULL


# ---------------------------------------------------------------------------
# 4. worker 清理核心：中断任务被重置（worker 重启/超时恢复）
# ---------------------------------------------------------------------------


def test_worker_cleanup_resets_interrupted_running_task():
    """模拟 worker 重启：running 超 20 分钟无更新的任务被重置为 failed，原因写明。"""
    repo, _, engine = _repository()

    t = _make_task(is_preview=True)  # 预览任务
    repo.create(t)
    t.mark_processing()  # running
    repo.update(t)
    _age_task(engine, t.id, updated_minutes=25)  # 25 分钟无进度更新

    cleaned = _startup.cleanup_stale_running_with_session(repo, 20)
    assert cleaned == 1

    saved = repo.get(t.id)
    assert saved.status == GenerationTaskStatus.FAILED
    assert "中断" in saved.error_message
    assert saved.error_info.get("error_type") == "WorkerInterrupted"
    assert saved.completed_at is not None


def test_worker_cleanup_keeps_healthy_running_task():
    """正常运行中（5 分钟前有更新）的任务不被误杀。"""
    repo, _, engine = _repository()

    t = _make_task()
    repo.create(t)
    t.mark_processing()
    repo.update(t)
    _age_task(engine, t.id, updated_minutes=5)

    assert _startup.cleanup_stale_running_with_session(repo, 20) == 0
    assert repo.get(t.id).status == GenerationTaskStatus.RUNNING


def test_worker_cleanup_resets_stale_pending_task():
    """卡 pending 超 15 分钟（worker 停止消费）的任务被重置，释放限流名额。"""
    repo, _, engine = _repository()

    t = _make_task(is_preview=True)
    repo.create(t)  # 一直 pending
    _age_task(engine, t.id, created_minutes=20)

    cleaned = _startup.cleanup_stale_pending_with_session(repo, 15)
    assert cleaned == 1

    saved = repo.get(t.id)
    assert saved.status == GenerationTaskStatus.FAILED
    assert saved.error_info.get("error_type") == "PendingTimeout"
    # 释放名额后 pending 计数归零，新请求不再被 429 误伤
    assert repo.count_pending_total() == 0


def test_worker_cleanup_pending_keeps_recent():
    """刚创建 3 分钟的 pending 任务不清理。"""
    repo, _, engine = _repository()

    t = _make_task()
    repo.create(t)
    _age_task(engine, t.id, created_minutes=3)

    assert _startup.cleanup_stale_pending_with_session(repo, 15) == 0
    assert repo.get(t.id).status == GenerationTaskStatus.PENDING


def test_worker_cleanup_multiple_orphans_all_reset():
    """3 个卡死 running 任务（工单实测：3 个预览卡 80% 超 10 小时）全部恢复。"""
    repo, _, engine = _repository()

    ids = []
    for i in range(3):
        t = _make_task(project_id=f"p{i}", is_preview=True)
        repo.create(t)
        t.mark_processing()
        repo.update(t)
        _age_task(engine, t.id, updated_minutes=600)  # 10 小时
        ids.append(t.id)

    cleaned = _startup.cleanup_stale_running_with_session(repo, 20)
    assert cleaned == 3
    for tid in ids:
        assert repo.get(tid).status == GenerationTaskStatus.FAILED