"""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