Files
zsglpt/tests/test_browser_pool_worker.py

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