160 lines
5.6 KiB
Python
160 lines
5.6 KiB
Python
import time
|
|
import threading
|
|
from app.routers.websites import _with_background_heartbeat, _try_start_task, _finish_task
|
|
import app.routers.websites as _websites_module
|
|
|
|
|
|
def test_heartbeat_emitted_during_blocking_operation():
|
|
"""生成器阻塞超过心跳间隔时，仍然能收到 heartbeat 事件。"""
|
|
# 模拟一个每 6 秒产生一条 item 的生成器
|
|
def slow_gen():
|
|
yield "item", {"n": 1}
|
|
time.sleep(7) # 阻塞时间 > interval=4
|
|
yield "item", {"n": 2}
|
|
yield "complete", {"success": True, "message": "done"}
|
|
|
|
events = list(_with_background_heartbeat(slow_gen, interval=4, timeout_factor=5))
|
|
|
|
heartbeat_count = sum(1 for e, _ in events if e == "heartbeat")
|
|
assert heartbeat_count >= 1, f"阻塞 7s 应产生至少 1 个 heartbeat，实际 {heartbeat_count}"
|
|
items = [(e, d) for e, d in events if e == "item"]
|
|
assert len(items) == 2
|
|
assert items[0][1]["n"] == 1
|
|
assert items[1][1]["n"] == 2
|
|
|
|
|
|
def test_heartbeat_multiple_during_long_block():
|
|
"""长时间阻塞（>2 个心跳周期）时产生多个 heartbeat。"""
|
|
def long_block_gen():
|
|
yield "item", {"n": 1}
|
|
time.sleep(11) # > 2 * interval=5
|
|
yield "complete", {"success": True, "message": "done"}
|
|
|
|
events = []
|
|
for event, data in _with_background_heartbeat(long_block_gen, interval=5, timeout_factor=5):
|
|
events.append(event)
|
|
if event == "complete":
|
|
break
|
|
|
|
heartbeat_count = sum(1 for e in events if e == "heartbeat")
|
|
assert heartbeat_count >= 2, f"11s 阻塞 (interval=5) 应产生至少 2 个 heartbeat，实际 {heartbeat_count}"
|
|
|
|
|
|
def test_queue_full_does_not_deadlock():
|
|
"""生产速度超过消费速度且队列满时，后台线程不会永久阻塞。"""
|
|
n_items = 300
|
|
|
|
def fast_gen():
|
|
for i in range(n_items):
|
|
yield "item", {"n": i}
|
|
yield "complete", {"success": True, "message": "done"}
|
|
|
|
# 使用大 timeout_factor 确保主线程有充足时间消费
|
|
events = list(_with_background_heartbeat(fast_gen, interval=10, timeout_factor=20))
|
|
|
|
item_count = sum(1 for e, _ in events if e == "item")
|
|
assert item_count == n_items, f"应有 {n_items} 个 item，实际 {item_count}"
|
|
assert any(e == "complete" for e, _ in events), "应收到 complete 事件"
|
|
|
|
|
|
def test_generator_close_stops_background_thread_immediately():
|
|
"""生成器关闭后后台线程停止，不阻塞 join 超时。"""
|
|
thread_done = threading.Event()
|
|
|
|
def infinite_gen():
|
|
try:
|
|
i = 0
|
|
while True:
|
|
yield "item", {"n": i}
|
|
i += 1
|
|
finally:
|
|
thread_done.set()
|
|
|
|
gen = _with_background_heartbeat(infinite_gen, interval=10, timeout_factor=10)
|
|
|
|
# 消费 3 个事件
|
|
count = 0
|
|
for event, data in gen:
|
|
count += 1
|
|
if count >= 3:
|
|
break
|
|
|
|
# 关闭生成器模拟客户端断开
|
|
gen.close()
|
|
|
|
# 后台线程应快速响应 cancel 并退出
|
|
assert thread_done.wait(timeout=10), "后台线程未在 10s 内停止"
|
|
|
|
|
|
def test_worker_exception_caught_as_error_event():
|
|
"""生成器内部抛出异常时，产生 error 事件而非崩溃。"""
|
|
def crashing_gen():
|
|
yield "item", {"n": 1}
|
|
raise RuntimeError("模拟意外错误")
|
|
|
|
events = list(_with_background_heartbeat(crashing_gen, interval=10, timeout_factor=5))
|
|
|
|
error_events = [(e, d) for e, d in events if e == "error"]
|
|
assert len(error_events) == 1
|
|
assert "模拟意外错误" in error_events[0][1]["message"]
|
|
# 不应收到 complete
|
|
assert not any(e == "complete" for e, _ in events)
|
|
|
|
|
|
def test_worker_timeout_emits_error_and_cancels():
|
|
"""worker 长时间不产生事件时，主线程超时并产生 error。"""
|
|
def stuck_gen():
|
|
yield "start", {}
|
|
time.sleep(15) # 远远超过 queue.get 超时时间
|
|
yield "complete", {"success": True, "message": "done"}
|
|
|
|
# interval=5, timeout_factor=0.5 → queue.get 超时 = 2.5s
|
|
# 心跳间隔(5s) > queue.get 超时(2.5s)，所以不会因心跳重置超时计数器
|
|
events = list(_with_background_heartbeat(stuck_gen, interval=5, timeout_factor=0.5))
|
|
|
|
error_events = [(e, d) for e, d in events if e == "error"]
|
|
assert len(error_events) == 1
|
|
assert "长时间无响应" in error_events[0][1]["message"]
|
|
|
|
|
|
def test_lock_held_until_worker_finishes_after_disconnect():
|
|
"""生成器关闭后任务锁保持持有，直到业务线程真正退出。"""
|
|
wid = 9999
|
|
|
|
ok, _task_token = _try_start_task(wid, "test_lock")
|
|
assert ok, "应成功获取锁"
|
|
|
|
worker_done = threading.Event()
|
|
|
|
def blocking_gen():
|
|
try:
|
|
yield "item", {"n": 1}
|
|
time.sleep(5) # 模拟长时间阻塞（HTTP 请求等）
|
|
yield "complete", {"success": True, "message": "done"}
|
|
finally:
|
|
worker_done.set()
|
|
|
|
gen = _with_background_heartbeat(
|
|
blocking_gen, interval=10, timeout_factor=10,
|
|
task_wid=wid, task_type="test_lock", task_token=_task_token,
|
|
)
|
|
|
|
# 消费一个 item 后关闭生成器（模拟客户端断开）
|
|
next(gen)
|
|
gen.close()
|
|
|
|
# 关闭后锁应当仍然被持有（worker 还在 sleep(5) 中）
|
|
ok, _ = _try_start_task(wid, "test_lock")
|
|
assert not ok, "断开后锁应继续保持持有"
|
|
|
|
# 等待 worker 真正结束
|
|
assert worker_done.wait(timeout=10), "worker 未在预期时间内结束"
|
|
|
|
# worker 结束后，_finish_task 已在 _run() 的 finally 中执行，
|
|
# 新请求应能获取锁
|
|
ok, _ = _try_start_task(wid, "test_lock")
|
|
assert ok, "worker 结束后应能重新获取锁"
|
|
|
|
# 清理
|
|
_finish_task(wid, "test_lock", _task_token)
|