53c78e8e3c
主要更新: - 新增 security/ 安全模块 (风险评估、威胁检测、蜜罐等) - Dockerfile 添加 curl 以支持 Docker 健康检查 - 前端页面更新 (管理后台、用户端) - 数据库迁移和 schema 更新 - 新增 kdocs 上传服务 - 添加安全相关测试用例 Co-Authored-By: Claude <noreply@anthropic.com>
75 lines
2.1 KiB
Python
75 lines
2.1 KiB
Python
from __future__ import annotations
|
|
|
|
import queue
|
|
|
|
from browser_pool_worker import BrowserWorker
|
|
|
|
|
|
class _AlwaysFailEnsureWorker(BrowserWorker):
|
|
def __init__(self, *, worker_id: int, task_queue: queue.Queue):
|
|
super().__init__(worker_id=worker_id, task_queue=task_queue, pre_warm=False)
|
|
self.ensure_calls = 0
|
|
|
|
def _ensure_browser(self) -> bool: # noqa: D401 - matching base naming
|
|
self.ensure_calls += 1
|
|
if self.ensure_calls >= 2:
|
|
self.running = False
|
|
return False
|
|
|
|
def _close_browser(self):
|
|
self.browser_instance = None
|
|
|
|
|
|
def test_requeue_task_when_browser_unavailable():
|
|
task_queue: queue.Queue = queue.Queue()
|
|
callback_calls: list[tuple[object, object]] = []
|
|
|
|
def callback(result, error):
|
|
callback_calls.append((result, error))
|
|
|
|
task = {
|
|
"func": lambda *_args, **_kwargs: None,
|
|
"args": (),
|
|
"kwargs": {},
|
|
"callback": callback,
|
|
"retry_count": 0,
|
|
}
|
|
|
|
worker = _AlwaysFailEnsureWorker(worker_id=1, task_queue=task_queue)
|
|
worker.start()
|
|
task_queue.put(task)
|
|
worker.join(timeout=5)
|
|
|
|
assert worker.is_alive() is False
|
|
assert worker.ensure_calls == 2 # 本地最多尝试2次创建执行环境
|
|
assert callback_calls == [] # 第一次失败会重新入队，不应立即回调失败
|
|
|
|
requeued = task_queue.get_nowait()
|
|
assert requeued["retry_count"] == 1
|
|
|
|
|
|
def test_fail_task_after_second_assignment():
|
|
task_queue: queue.Queue = queue.Queue()
|
|
callback_calls: list[tuple[object, object]] = []
|
|
|
|
def callback(result, error):
|
|
callback_calls.append((result, error))
|
|
|
|
task = {
|
|
"func": lambda *_args, **_kwargs: None,
|
|
"args": (),
|
|
"kwargs": {},
|
|
"callback": callback,
|
|
"retry_count": 1, # 已重新分配过1次
|
|
}
|
|
|
|
worker = _AlwaysFailEnsureWorker(worker_id=1, task_queue=task_queue)
|
|
worker.start()
|
|
task_queue.put(task)
|
|
worker.join(timeout=5)
|
|
|
|
assert worker.is_alive() is False
|
|
assert callback_calls == [(None, "执行环境不可用")]
|
|
assert worker.total_tasks == 1
|
|
assert worker.failed_tasks == 1
|