xiaoxia-saas/tests/unit/test_celery_queue_isolation_1714.py

"""#1714 队列隔离 + 作废消息清除 单元测试。

覆盖：
1. task_routes：generate_video → generation，ingest_asset/classify/duplication → transcode
2. purge_stale_messages_from_queues：Redis 队列中作废任务消息被物理移除，未命中保留
3. revoke_and_purge：revoke 广播 + 队列清理同时生效
4. ensure_task_claimable：终态任务抛 StaleTaskDiscarded，pending 放行
"""

from __future__ import annotations

from unittest.mock import MagicMock

import pytest
from celery import Celery

from packages.shared.celery_orphan_guard import (
    StaleTaskDiscarded,
    _extract_business_ids,
    ensure_task_claimable,
    purge_stale_messages_from_queues,
    revoke_and_purge,
)
from packages.shared.celery_queues import (
    QUEUE_GENERATION,
    QUEUE_TRANSCODE,
    apply_queue_settings,
    task_routes,
)

BROKER_URL = "redis://localhost:6379/15"
TEST_QUEUES = ("_test_gen_q", "_test_transcode_q")


# ── 1. 路由表 ──────────────────────────────────────────────────────────


def test_routes_send_generation_to_generation_queue():
    assert task_routes["worker.generate_video"]["queue"] == QUEUE_GENERATION


def test_routes_send_ingest_to_transcode_queue():
    assert task_routes["worker.ingest_asset"]["queue"] == QUEUE_TRANSCODE
    assert task_routes["worker.classify_asset"]["queue"] == QUEUE_TRANSCODE
    assert task_routes["worker.process_duplication_check"]["queue"] == QUEUE_TRANSCODE
    assert task_routes["worker.check_duplicate"]["queue"] == QUEUE_TRANSCODE


def test_apply_queue_settings_configures_celery_app():
    app = Celery("test-routes")
    apply_queue_settings(app)
    queue_names = {q.name for q in app.conf.task_queues}
    assert queue_names == {"generation", "transcode", "celery"}
    assert app.conf.task_default_queue == "celery"


# ── Redis 队列消息清理（需要本地 redis；不可用时 skip） ─────────────────


def _redis_available() -> bool:
    try:
        import redis

        return bool(redis.Redis.from_url(BROKER_URL).ping())
    except Exception:
        return False


@pytest.fixture()
def redis_client():
    import redis

    client = redis.Redis.from_url(BROKER_URL)
    for q in TEST_QUEUES:
        client.delete(q)
    yield client
    for q in TEST_QUEUES:
        client.delete(q)


def _publish(app: Celery, queue: str, celery_id: str, business_id: str) -> None:
    from kombu import Queue
    from kombu.pools import producers

    with app.connection_for_write() as conn:
        with producers[conn].acquire(block=True) as prod:
            prod.publish(
                (business_id,),
                exchange="",
                routing_key=queue,
                serializer="json",
                headers={"id": celery_id, "task": "worker.generate_video"},
                retry=False,
                delivery_mode=1,
                declare=[Queue(queue, routing_key=queue, durable=False)],
            )


@pytest.mark.skipif(not _redis_available(), reason="本地 redis 不可用")
def test_purge_removes_stale_business_message_and_keeps_others(redis_client):
    app = Celery("test-purge")
    app.conf.broker_url = BROKER_URL
    _publish(app, TEST_QUEUES[0], "celery-1", "task-KEEP-A")
    _publish(app, TEST_QUEUES[0], "celery-2", "task-STALE-B")
    _publish(app, TEST_QUEUES[0], "celery-3", "task-KEEP-C")
    _publish(app, TEST_QUEUES[1], "celery-4", "task-STALE-B")  # 同一业务任务在转码队列？不应出现但验证全队列扫描

    removed = purge_stale_messages_from_queues(BROKER_URL, TEST_QUEUES, business_task_ids={"task-STALE-B"})
    assert removed == 2

    remaining = []
    for raw in redis_client.lrange(TEST_QUEUES[0], 0, -1):
        _celery_id, biz_id = _extract_business_ids(raw)
        remaining.append(biz_id)
    assert set(remaining) == {"task-KEEP-A", "task-KEEP-C"}
    assert redis_client.llen(TEST_QUEUES[1]) == 0


@pytest.mark.skipif(not _redis_available(), reason="本地 redis 不可用")
def test_purge_matches_by_celery_message_id(redis_client):
    app = Celery("test-purge-msg-id")
    app.conf.broker_url = BROKER_URL
    _publish(app, TEST_QUEUES[0], "celery-stale-id", "task-X")
    _publish(app, TEST_QUEUES[0], "celery-good-id", "task-Y")

    removed = purge_stale_messages_from_queues(BROKER_URL, TEST_QUEUES, celery_task_ids={"celery-stale-id"})
    assert removed == 1
    assert redis_client.llen(TEST_QUEUES[0]) == 1


@pytest.mark.skipif(not _redis_available(), reason="本地 redis 不可用")
def test_revoke_and_purge_calls_control_revoke(redis_client):
    app = Celery("test-revoke")
    app.conf.broker_url = BROKER_URL
    app.control = MagicMock()
    _publish(app, TEST_QUEUES[0], "celery-revoke-1", "task-R")

    removed = revoke_and_purge(
        app,
        BROKER_URL,
        business_task_ids={"task-R"},
        celery_task_ids={"celery-revoke-1"},
        queue_names=TEST_QUEUES,
    )
    assert removed == 1
    app.control.revoke.assert_called_once_with("celery-revoke-1")


def test_purge_empty_ids_is_noop():
    assert purge_stale_messages_from_queues(BROKER_URL, TEST_QUEUES) == 0


# ── 2. 执行前状态守卫 ──────────────────────────────────────────────────


def test_guard_allows_pending():
    status = ensure_task_claimable("t1", lambda _id: "pending", task_label="generation")
    assert status == "pending"


def test_guard_rejects_failed():
    with pytest.raises(StaleTaskDiscarded) as exc:
        ensure_task_claimable("t2", lambda _id: "failed", task_label="generation")
    assert exc.value.task_id == "t2"
    assert exc.value.status == "failed"


def test_guard_rejects_cancelled_and_completed():
    with pytest.raises(StaleTaskDiscarded):
        ensure_task_claimable("t3", lambda _id: "cancelled")
    with pytest.raises(StaleTaskDiscarded):
        ensure_task_claimable("t4", lambda _id: "completed")


def test_guard_missing_task_returns_empty():
    assert ensure_task_claimable("t5", lambda _id: None) == ""