xiaoxia-saas/tests/unit/test_task_enqueue.py

"""task_enqueue 单测 — 队列限流 + 安全入队逻辑."""

from __future__ import annotations

from unittest.mock import MagicMock, patch

import pytest
from app.core.task_enqueue import (
    GLOBAL_PENDING_LIMIT,
    USER_PENDING_LIMIT,
    GlobalQueueFull,
    UserPendingLimitExceeded,
    check_queue_limits,
    safe_enqueue_generation_task,
)

# ── Fixtures / Helpers ─────────────────────────────────────────────────────


class MockRepository:
    """Mock 任务仓储，用计数器模拟 pending 数量."""

    def __init__(self, global_count: int = 0, user_count: int = 0):
        self._global = global_count
        self._user = user_count
        self.update_called = 0

    def count_pending_total(self) -> int:
        return self._global

    def count_pending_by_user(self, user_id: str) -> int:
        return self._user

    def update(self, task):
        self.update_called += 1


def make_mock_task(task_id: str = "task-1"):
    task = MagicMock()
    task.id = task_id
    task.status = "pending"
    task.mark_failed = MagicMock()
    return task


# ── check_queue_limits ────────────────────────────────────────────────────


class TestCheckQueueLimits:
    """check_queue_limits 预检查限流."""

    def test_below_limits_passes(self):
        repo = MockRepository(global_count=5, user_count=1)
        # 不抛异常就是通过
        check_queue_limits("user-1", repo)

    def test_global_at_limit_raises(self):
        """达到全局上限即拒绝."""
        repo = MockRepository(global_count=GLOBAL_PENDING_LIMIT, user_count=1)
        with pytest.raises(GlobalQueueFull) as exc_info:
            check_queue_limits("user-1", repo)
        assert exc_info.value.pending_count == GLOBAL_PENDING_LIMIT
        assert exc_info.value.limit == GLOBAL_PENDING_LIMIT

    def test_global_over_limit_raises(self):
        repo = MockRepository(global_count=GLOBAL_PENDING_LIMIT + 1, user_count=1)
        with pytest.raises(GlobalQueueFull):
            check_queue_limits("user-1", repo)

    def test_user_at_limit_raises(self):
        """达到用户上限即拒绝."""
        repo = MockRepository(global_count=5, user_count=USER_PENDING_LIMIT)
        with pytest.raises(UserPendingLimitExceeded) as exc_info:
            check_queue_limits("user-1", repo)
        assert exc_info.value.user_id == "user-1"
        assert exc_info.value.pending_count == USER_PENDING_LIMIT
        assert exc_info.value.limit == USER_PENDING_LIMIT

    def test_user_over_limit_raises(self):
        repo = MockRepository(global_count=5, user_count=USER_PENDING_LIMIT + 1)
        with pytest.raises(UserPendingLimitExceeded):
            check_queue_limits("user-1", repo)

    def test_global_priority_over_user(self):
        """全局和用户都超限时，优先抛全局异常."""
        repo = MockRepository(
            global_count=GLOBAL_PENDING_LIMIT + 1,
            user_count=USER_PENDING_LIMIT + 1,
        )
        with pytest.raises(GlobalQueueFull):
            check_queue_limits("user-1", repo)

    def test_empty_user_id_skips_user_check(self):
        """user_id 为空时跳过用户级检查."""
        repo = MockRepository(global_count=5, user_count=999)
        # 不抛异常 = 通过（只检查全局）
        check_queue_limits("", repo)

    def test_custom_limits(self):
        """支持自定义限流阈值."""
        repo = MockRepository(global_count=5, user_count=5)
        # 默认阈值下 user 5 > 3 会被拒
        with pytest.raises(UserPendingLimitExceeded):
            check_queue_limits("u1", repo)

        # 自定义更高阈值就能通过
        check_queue_limits("u1", repo, user_pending_limit=10, global_pending_limit=10)


# ── safe_enqueue_generation_task ──────────────────────────────────────────


class TestSafeEnqueueGenerationTask:
    """safe_enqueue_generation_task 安全入队."""

    @patch("app.core.task_enqueue.celery_app")
    def test_success_path(self, mock_celery):
        """正常路径：入队前检查通过 → 发送Celery → 入队后检查通过."""
        repo = MockRepository(global_count=1, user_count=1)
        task = make_mock_task()

        result = safe_enqueue_generation_task(task, repo, user_id="user-1")

        assert result is True
        mock_celery.send_task.assert_called_once_with("worker.generate_video", args=[task.id])
        task.mark_failed.assert_not_called()

    @patch("app.core.task_enqueue.celery_app")
    def test_no_user_id_skips_user_check(self, mock_celery):
        """不传 user_id 跳过用户级限流."""
        repo = MockRepository(global_count=1, user_count=999)
        task = make_mock_task()

        result = safe_enqueue_generation_task(task, repo, user_id="")
        assert result is True

    @patch("app.core.task_enqueue.celery_app")
    def test_precheck_global_over_marks_failed(self, mock_celery):
        """入队前全局超限：标记 failed，抛异常."""
        repo = MockRepository(global_count=GLOBAL_PENDING_LIMIT + 1, user_count=0)
        task = make_mock_task()

        with pytest.raises(GlobalQueueFull):
            safe_enqueue_generation_task(task, repo, user_id="user-1")

        task.mark_failed.assert_called_once()
        mock_celery.send_task.assert_not_called()
        assert repo.update_called == 1

    @patch("app.core.task_enqueue.celery_app")
    def test_precheck_user_over_marks_failed(self, mock_celery):
        """入队前用户超限：标记 failed，抛异常."""
        repo = MockRepository(global_count=5, user_count=USER_PENDING_LIMIT + 1)
        task = make_mock_task()

        with pytest.raises(UserPendingLimitExceeded):
            safe_enqueue_generation_task(task, repo, user_id="user-1")

        task.mark_failed.assert_called_once()
        mock_celery.send_task.assert_not_called()

    @patch("app.core.task_enqueue.celery_app")
    def test_celery_send_false_returns_false(self, mock_celery):
        """Celery 发送失败：返回 False，任务标记 failed."""
        repo = MockRepository(global_count=1, user_count=1)
        task = make_mock_task()
        mock_celery.send_task.side_effect = Exception("celery down")

        result = safe_enqueue_generation_task(task, repo, user_id="user-1")

        assert result is False
        task.mark_failed.assert_called_once()
        assert "入队失败" in task.mark_failed.call_args[0][0]

    @patch("app.core.task_enqueue.celery_app")
    def test_celery_send_failure_update_also_fails(self, mock_celery):
        """Celery 发送失败 + mark_failed 更新也失败：不崩溃."""
        repo = MockRepository(global_count=1, user_count=1)
        repo.update = MagicMock(side_effect=Exception("db down"))
        task = make_mock_task()
        mock_celery.send_task.side_effect = Exception("celery down")

        result = safe_enqueue_generation_task(task, repo, user_id="user-1")

        assert result is False
        # 不抛异常就是胜利

    @patch("app.core.task_enqueue.celery_app")
    def test_postcheck_global_over_rollback(self, mock_celery):
        """入队后全局超限（并发竞态）：回滚标记 failed，抛异常."""
        # 入队前刚好通过，但入队后再查发现超限
        call_count = [0]

        def count_pending_total_side_effect():
            call_count[0] += 1
            if call_count[0] == 1:  # 入队前检查
                return GLOBAL_PENDING_LIMIT  # 等于上限，用 > 判断所以通过
            return GLOBAL_PENDING_LIMIT + 1  # 入队后再查，超限

        repo = MockRepository(global_count=GLOBAL_PENDING_LIMIT, user_count=0)
        repo.count_pending_total = MagicMock(side_effect=count_pending_total_side_effect)
        task = make_mock_task()

        with pytest.raises(GlobalQueueFull):
            safe_enqueue_generation_task(task, repo, user_id="user-1")

        # 异常是 GlobalQueueFull 类型，且任务已被标记为 failed（含"入队后"原因）
        task.mark_failed.assert_called_once()
        assert "入队后" in task.mark_failed.call_args[0][0]
        mock_celery.send_task.assert_called_once()

    @patch("app.core.task_enqueue.celery_app")
    def test_postcheck_user_over_rollback(self, mock_celery):
        """入队后用户超限：回滚标记 failed，抛异常."""
        repo = MockRepository(global_count=5, user_count=USER_PENDING_LIMIT)
        # 入队前用 > 判断，等于上限通过；入队后模拟并发超限
        original_user_count = repo.count_pending_by_user

        call_count = [0]

        def count_by_user_side_effect(user_id):
            call_count[0] += 1
            if call_count[0] <= 1:  # 入队前
                return USER_PENDING_LIMIT  # 用 > 判断，等于时通过
            return USER_PENDING_LIMIT + 1  # 入队后，超限

        repo.count_pending_by_user = MagicMock(side_effect=count_by_user_side_effect)
        task = make_mock_task()

        with pytest.raises(UserPendingLimitExceeded):
            safe_enqueue_generation_task(task, repo, user_id="user-1")

        task.mark_failed.assert_called_once()

    @patch("app.core.task_enqueue.celery_app")
    def test_log_task_status_enabled(self, mock_celery):
        """log_task_status=True 时日志中包含状态."""
        repo = MockRepository(global_count=1, user_count=1)
        task = make_mock_task()

        result = safe_enqueue_generation_task(task, repo, user_id="user-1", log_task_status=True)
        assert result is True

    @patch("app.core.task_enqueue.celery_app")
    def test_custom_limits_in_enqueue(self, mock_celery):
        """自定义限流阈值用于入队检查."""
        repo = MockRepository(global_count=5, user_count=5)
        task = make_mock_task()

        # 默认阈值下用户 5 > 3 会被拒
        with pytest.raises(UserPendingLimitExceeded):
            safe_enqueue_generation_task(task, repo, user_id="user-1")

        # 重置 mock 计数
        task.mark_failed.reset_mock()

        # 调大阈值后通过
        result = safe_enqueue_generation_task(
            task,
            repo,
            user_id="user-1",
            user_pending_limit=10,
            global_pending_limit=10,
        )
        assert result is True


# ── 异常类 ────────────────────────────────────────────────────────────────


class TestExceptionClasses:
    """异常类消息格式."""

    def test_user_pending_limit_message(self):
        exc = UserPendingLimitExceeded("u1", 5, 3)
        assert "u1" in str(exc)
        assert "5" in str(exc)
        assert "3" in str(exc)

    def test_global_queue_full_message(self):
        exc = GlobalQueueFull(25, 20)
        assert "25" in str(exc)
        assert "20" in str(exc)