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