"""Issue #1714：孤儿/超时清理标记 failed 时必须撤销并清除 Redis 队列消息。 覆盖： - cleanup_stale_pending_with_session_ids：超时 pending 标记 failed 并返回 (task_id, celery_task_id)，worker 清理流程据此 revoke + purge 队列消息 - 队列中对应业务任务的 celery 消息被物理移除（作废消息不会重投执行） - 旧仓储（无 _with_ids 方法）降级为计数模式，不抛异常 - cleanup_stale_running_with_ids 同样返回 id 列表 """ from __future__ import annotations import sys from datetime import UTC, datetime, timedelta, timezone from pathlib import Path from unittest.mock import MagicMock import pytest from sqlalchemy import create_engine, text from sqlalchemy.orm import sessionmaker sys.path.insert(0, str(Path(__file__).resolve().parents[2] / "apps" / "api")) 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 BROKER_URL = "redis://localhost:6379/15" TEST_QUEUE = "_test_revoke_q" def _repository(): engine = create_engine("sqlite:///:memory:") 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 _redis_available() -> bool: try: import redis return bool(redis.Redis.from_url(BROKER_URL).ping()) except Exception: return False # ── 仓储层：返回 ids ──────────────────────────────────────────────────── def test_cleanup_stale_pending_returns_ids_with_celery_task_id(): repo, _, engine = _repository() task = _make_task() task.celery_task_id = "celery-msg-id-001" repo.create(task) # created_at 改到 60 分钟前 with engine.connect() as conn: conn.execute( text("UPDATE generation_tasks SET created_at = :ts WHERE id = :id"), {"ts": datetime.now(UTC) - timedelta(minutes=60), "id": task.id}, ) conn.commit() items = repo.cleanup_stale_pending_with_ids(timeout_minutes=45) assert len(items) == 1 biz_id, celery_id = items[0] assert biz_id == task.id assert celery_id == "celery-msg-id-001" saved = repo.get(task.id) assert saved.status == GenerationTaskStatus.FAILED def test_cleanup_stale_running_returns_ids(): repo, _, engine = _repository() task = _make_task() repo.create(task) task.mark_processing() task.celery_task_id = "celery-msg-id-002" repo.update(task) with engine.connect() as conn: conn.execute( text("UPDATE generation_tasks SET updated_at = :ts WHERE id = :id"), {"ts": datetime.now(UTC) - timedelta(minutes=60), "id": task.id}, ) conn.commit() items = repo.cleanup_stale_running_with_ids(timeout_minutes=20) assert len(items) == 1 assert items[0][0] == task.id assert items[0][1] == "celery-msg-id-002" assert repo.get(task.id).status == GenerationTaskStatus.FAILED def test_legacy_repo_without_with_ids_falls_back_to_count(): """旧仓储只有 cleanup_stale_pending（返回 int）时降级可用，不抛异常。""" # worker 模块加载（标准 mock 模式） saved = set(sys.modules.keys()) mock_db = MagicMock() mock_db.SessionLocal = MagicMock() sys.modules["worker_app.db"] = mock_db sys.modules["worker_app.core.config"] = MagicMock() mock_celery = MagicMock() mock_celery.celery_app.task = MagicMock( side_effect=(lambda *a, **k: (a[0] if a and callable(a[0]) else (lambda f: f))) ) sys.modules["worker_app.celery_app"] = mock_celery worker_path = str(Path(__file__).resolve().parents[2] / "apps" / "worker") if worker_path not in sys.path: sys.path.insert(0, worker_path) from worker_app.tasks import _startup # noqa: E402 class LegacyRepo: def cleanup_stale_pending(self, timeout_minutes): # noqa: ARG002 return 3 def cleanup_stale_running(self, timeout_minutes): # noqa: ARG002 return 2 items_p = _startup.cleanup_stale_pending_with_session_ids(LegacyRepo(), 45) items_r = _startup.cleanup_stale_running_with_session_ids(LegacyRepo(), 20) assert len(items_p) == 3 assert len(items_r) == 2 for key in list(sys.modules.keys()): if key not in saved and not key.startswith("video_processing"): del sys.modules[key] # ── 端到端：清理 → 队列消息被移除（作废消息不重投） ──────────────────── @pytest.mark.skipif(not _redis_available(), reason="本地 redis 不可用") def test_stale_pending_cleanup_purges_redis_message(): """任务标 failed 后，其在 Redis 队列里的 celery 消息被清除，不会被重投。""" import redis from celery import Celery from kombu import Queue from kombu.pools import producers from packages.shared.celery_orphan_guard import purge_stale_messages_from_queues repo, _, engine = _repository() task = _make_task() task.celery_task_id = "celery-stale-xyz" repo.create(task) with engine.connect() as conn: conn.execute( text("UPDATE generation_tasks SET created_at = :ts WHERE id = :id"), {"ts": datetime.now(UTC) - timedelta(minutes=60), "id": task.id}, ) conn.commit() # 模拟该任务的 celery 消息仍在 generation 队列里（worker 下线期间未消费） client = redis.Redis.from_url(BROKER_URL) client.delete(TEST_QUEUE) app = Celery("test-e2e-revoke") app.conf.broker_url = BROKER_URL with app.connection_for_write() as conn: with producers[conn].acquire(block=True) as prod: # 作废任务消息 prod.publish( (task.id,), exchange="", routing_key=TEST_QUEUE, serializer="json", headers={"id": "celery-stale-xyz", "task": "worker.generate_video"}, retry=False, delivery_mode=1, declare=[Queue(TEST_QUEUE, routing_key=TEST_QUEUE, durable=False)], ) # 另一条正常任务消息（必须保留） prod.publish( ("other-task-id",), exchange="", routing_key=TEST_QUEUE, serializer="json", headers={"id": "celery-keep", "task": "worker.generate_video"}, retry=False, delivery_mode=1, ) assert client.llen(TEST_QUEUE) == 2 # 执行清理（与 worker beat 相同流程：标 failed → 拿 ids → purge） items = repo.cleanup_stale_pending_with_ids(timeout_minutes=45) biz_ids = [bid for bid, _ in items] celery_ids = [cid for _, cid in items if cid] removed = purge_stale_messages_from_queues( BROKER_URL, (TEST_QUEUE,), business_task_ids=biz_ids, celery_task_ids=celery_ids ) assert removed == 1 assert client.llen(TEST_QUEUE) == 1 # 正常任务消息保留 client.delete(TEST_QUEUE)