fix: stream long-running website operations

backend/test_background_heartbeat.py

@@ -0,0 +1,159 @@
+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)