347 lines
11 KiB
Python
347 lines
11 KiB
Python
 
|
"""任务槽（hexi/core/task_slots.py）单元测试
|
|||
|
|
|
|||
|
|
只测注册表本身：执行/去重/进度/取消/推送/回收。模块只依赖 asyncio + nonebot.log，
|
|||
|
|
按路径裸加载即可（与 test_rate_limit 同法）。
|
|||
|
|
|
|||
|
|
任务体跑在线程池里，进度回调是"切回事件循环再改状态"，所以断言进度时必须等
|
|||
|
|
状态落地 —— 用 `_until()` 轮询事件循环，别用 `sleep` 猜时序。
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
import asyncio
|
|||
|
|
import importlib.util
|
|||
|
|
import sys
|
|||
|
|
import threading
|
|||
|
|
import time
|
|||
|
|
from pathlib import Path
|
|||
|
|
|
|||
|
|
import pytest
|
|||
|
|
|
|||
|
|
_MODULE_PATH = Path(__file__).resolve().parents[1] / "hexi" / "core" / "task_slots.py"
|
|||
|
|
_MOD_NAME = "hexi_core_task_slots_under_test"
|
|||
|
|
|
|||
|
|
|
|||
|
|
@pytest.fixture(scope="module")
|
|||
|
|
def ts():
|
|||
|
|
spec = importlib.util.spec_from_file_location(_MOD_NAME, _MODULE_PATH)
|
|||
|
|
module = importlib.util.module_from_spec(spec)
|
|||
|
|
sys.modules[_MOD_NAME] = module
|
|||
|
|
spec.loader.exec_module(module)
|
|||
|
|
yield module
|
|||
|
|
sys.modules.pop(_MOD_NAME, None)
|
|||
|
|
|
|||
|
|
|
|||
|
|
@pytest.fixture
|
|||
|
|
def reg(ts):
|
|||
|
|
"""每条用例一个干净的注册表（不用全局单例，避免用例互相污染）。"""
|
|||
|
|
return ts.TaskSlots()
|
|||
|
|
|
|||
|
|
|
|||
|
|
async def _until(pred, timeout: float = 2.0) -> None:
|
|||
|
|
"""轮询事件循环直到条件成立（进度是跨线程/跨回调落地的，不能靠猜时序）。"""
|
|||
|
|
deadline = time.monotonic() + timeout
|
|||
|
|
while time.monotonic() < deadline:
|
|||
|
|
if pred():
|
|||
|
|
return
|
|||
|
|
await asyncio.sleep(0.01)
|
|||
|
|
raise AssertionError("等待条件超时")
|
|||
|
|
|
|||
|
|
|
|||
|
|
# ───────────────────────── 执行与结果 ─────────────────────────
|
|||
|
|
|
|||
|
|
|
|||
|
|
async def test_runs_in_thread_and_records_result(reg):
|
|||
|
|
def work(report):
|
|||
|
|
return "done!"
|
|||
|
|
|
|||
|
|
task, reused = reg.submit(slot="plug:job", title="干活", func=work, plugin_id="plug")
|
|||
|
|
assert reused is False
|
|||
|
|
assert task.active # submit 即返回，此时后台任务还没被调度
|
|||
|
|
|
|||
|
|
await asyncio.wait_for(reg.wait(task), 2)
|
|||
|
|
assert task.status == "done"
|
|||
|
|
assert task.result == "done!"
|
|||
|
|
assert task.started_at > 0
|
|||
|
|
assert task.finished_at >= task.started_at
|
|||
|
|
assert task.duration >= 0
|
|||
|
|
assert reg.active("plug:job") is None
|
|||
|
|
|
|||
|
|
|
|||
|
|
async def test_exception_marks_failed_without_raising(reg):
|
|||
|
|
def boom(report):
|
|||
|
|
raise RuntimeError("炸了")
|
|||
|
|
|
|||
|
|
task, _ = reg.submit(slot="plug:boom", title="会炸", func=boom)
|
|||
|
|
await asyncio.wait_for(reg.wait(task), 2)
|
|||
|
|
assert task.status == "failed"
|
|||
|
|
assert "炸了" in task.error
|
|||
|
|
assert task.percent == 100 # 终态一律显示满（前端别再转圈）
|
|||
|
|
|
|||
|
|
|
|||
|
|
# ───────────────────────── 同槽去重 ─────────────────────────
|
|||
|
|
|
|||
|
|
|
|||
|
|
async def test_same_slot_reuses_running_task(reg):
|
|||
|
|
calls = []
|
|||
|
|
gate = threading.Event()
|
|||
|
|
|
|||
|
|
def slow(report):
|
|||
|
|
calls.append(1)
|
|||
|
|
gate.wait(2)
|
|||
|
|
return "ok"
|
|||
|
|
|
|||
|
|
first, reused = reg.submit(slot="plug:clean", title="清理", func=slow)
|
|||
|
|
second, reused2 = reg.submit(slot="plug:clean", title="清理（再点一次）", func=slow)
|
|||
|
|
|
|||
|
|
assert reused is False and reused2 is True
|
|||
|
|
assert second.id == first.id # 返回的是正在跑的那个
|
|||
|
|
gate.set()
|
|||
|
|
await asyncio.wait_for(reg.wait(first), 2)
|
|||
|
|
assert calls == [1] # 任务体只跑了一次
|
|||
|
|
|
|||
|
|
|
|||
|
|
async def test_same_slot_available_again_after_finish(reg):
|
|||
|
|
def work(report):
|
|||
|
|
return 1
|
|||
|
|
|
|||
|
|
first, _ = reg.submit(slot="plug:clean", title="第一次", func=work)
|
|||
|
|
await asyncio.wait_for(reg.wait(first), 2)
|
|||
|
|
|
|||
|
|
second, reused = reg.submit(slot="plug:clean", title="第二次", func=work)
|
|||
|
|
assert reused is False and second.id != first.id
|
|||
|
|
await asyncio.wait_for(reg.wait(second), 2)
|
|||
|
|
|
|||
|
|
|
|||
|
|
async def test_different_slots_run_in_parallel(reg):
|
|||
|
|
gate = threading.Event()
|
|||
|
|
|
|||
|
|
def slow(report):
|
|||
|
|
gate.wait(2)
|
|||
|
|
|
|||
|
|
a, _ = reg.submit(slot="plug:a", title="A", func=slow)
|
|||
|
|
b, reused = reg.submit(slot="plug:b", title="B", func=slow)
|
|||
|
|
assert reused is False and b.id != a.id
|
|||
|
|
gate.set()
|
|||
|
|
await asyncio.gather(reg.wait(a), reg.wait(b))
|
|||
|
|
|
|||
|
|
|
|||
|
|
# ───────────────────────── 进度 ─────────────────────────
|
|||
|
|
|
|||
|
|
|
|||
|
|
async def test_progress_fields_and_percent(reg):
|
|||
|
|
gate = threading.Event()
|
|||
|
|
|
|||
|
|
def work(report):
|
|||
|
|
report(done=3, total=12, phase="清理中", unit="个文件", message="已删除 2 个")
|
|||
|
|
gate.wait(2)
|
|||
|
|
return "ok"
|
|||
|
|
|
|||
|
|
task, _ = reg.submit(slot="plug:job", title="进度", func=work)
|
|||
|
|
await _until(lambda: task.done == 3)
|
|||
|
|
|
|||
|
|
assert (task.total, task.unit, task.phase) == (12, "个文件", "清理中")
|
|||
|
|
assert task.message == "已删除 2 个"
|
|||
|
|
assert task.percent == 25
|
|||
|
|
|
|||
|
|
payload = task.to_dict()["progress"]
|
|||
|
|
assert payload == {
|
|||
|
|
"percent": 25, "done": 3, "total": 12, "unit": "个文件",
|
|||
|
|
"phase": "清理中", "message": "已删除 2 个",
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
gate.set()
|
|||
|
|
await asyncio.wait_for(reg.wait(task), 2)
|
|||
|
|
|
|||
|
|
|
|||
|
|
async def test_unknown_total_reports_none_percent(reg):
|
|||
|
|
"""扫描阶段总数未知 → percent None（前端显示不确定进度而不是 0%）。"""
|
|||
|
|
gate = threading.Event()
|
|||
|
|
|
|||
|
|
def work(report):
|
|||
|
|
report(phase="扫描中", unit="个文件", message="正在统计文件…")
|
|||
|
|
gate.wait(2)
|
|||
|
|
|
|||
|
|
task, _ = reg.submit(slot="plug:job", title="扫描", func=work)
|
|||
|
|
await _until(lambda: task.phase == "扫描中")
|
|||
|
|
assert task.percent is None
|
|||
|
|
gate.set()
|
|||
|
|
await asyncio.wait_for(reg.wait(task), 2)
|
|||
|
|
|
|||
|
|
|
|||
|
|
async def test_warnings_accumulate(reg):
|
|||
|
|
def work(report):
|
|||
|
|
report(warning="跳过 a.mp4（被占用）")
|
|||
|
|
report(warning="跳过 b.mp4（被占用）")
|
|||
|
|
return (0, 2)
|
|||
|
|
|
|||
|
|
task, _ = reg.submit(slot="plug:job", title="部分失败", func=work)
|
|||
|
|
await asyncio.wait_for(reg.wait(task), 2)
|
|||
|
|
assert task.to_dict()["warnings"] == ["跳过 a.mp4（被占用）", "跳过 b.mp4（被占用）"]
|
|||
|
|
|
|||
|
|
|
|||
|
|
async def test_progress_is_throttled(reg):
|
|||
|
|
"""上千次上报不能变成上千次推送（按百分比桶 + 只在有意义的字段变化时发）。"""
|
|||
|
|
queue = reg.subscribe()
|
|||
|
|
count = 500
|
|||
|
|
|
|||
|
|
def work(report):
|
|||
|
|
for i in range(1, count + 1):
|
|||
|
|
report(done=i, total=count)
|
|||
|
|
return "ok"
|
|||
|
|
|
|||
|
|
task, _ = reg.submit(slot="plug:job", title="节流", func=work)
|
|||
|
|
await asyncio.wait_for(reg.wait(task), 2)
|
|||
|
|
|
|||
|
|
frames = 0
|
|||
|
|
while not queue.empty():
|
|||
|
|
queue.get_nowait()
|
|||
|
|
frames += 1
|
|||
|
|
# 理论上限：每个百分点一帧 + 起止几帧；远小于 500 次上报
|
|||
|
|
assert frames <= 110, f"推送 {frames} 帧，节流没生效"
|
|||
|
|
assert frames >= 2 # 但也不能一帧都不发
|
|||
|
|
reg.unsubscribe(queue)
|
|||
|
|
|
|||
|
|
|
|||
|
|
# ───────────────────────── 协作式取消 ─────────────────────────
|
|||
|
|
|
|||
|
|
|
|||
|
|
async def test_cancel_stops_cooperative_task(reg):
|
|||
|
|
def work(report):
|
|||
|
|
for i in range(1, 1000):
|
|||
|
|
if not report(done=i, total=999):
|
|||
|
|
return ("canceled", i) # 任务体自己收尾，带上已有进度
|
|||
|
|
time.sleep(0.005)
|
|||
|
|
return ("finished", 999)
|
|||
|
|
|
|||
|
|
task, _ = reg.submit(slot="plug:job", title="可取消", func=work)
|
|||
|
|
await _until(lambda: task.done >= 1)
|
|||
|
|
|
|||
|
|
assert reg.cancel(task.id) is True
|
|||
|
|
assert task.cancel_requested is True
|
|||
|
|
|
|||
|
|
await asyncio.wait_for(reg.wait(task), 2)
|
|||
|
|
assert task.status == "canceled"
|
|||
|
|
assert task.result[0] == "canceled"
|
|||
|
|
# 取消不是硬杀：任务体返回前的进度照常保留
|
|||
|
|
assert task.done < 999
|
|||
|
|
|
|||
|
|
|
|||
|
|
async def test_cancel_finished_or_unknown_task_is_noop(reg):
|
|||
|
|
def work(report):
|
|||
|
|
return 1
|
|||
|
|
|
|||
|
|
task, _ = reg.submit(slot="plug:job", title="快活", func=work)
|
|||
|
|
await asyncio.wait_for(reg.wait(task), 2)
|
|||
|
|
assert reg.cancel(task.id) is False
|
|||
|
|
assert reg.cancel("不存在的 id") is False
|
|||
|
|
|
|||
|
|
|
|||
|
|
# ───────────────────────── 订阅 / 回收 ─────────────────────────
|
|||
|
|
|
|||
|
|
|
|||
|
|
async def test_subscribe_receives_snapshots(reg):
|
|||
|
|
queue = reg.subscribe()
|
|||
|
|
assert queue.empty()
|
|||
|
|
|
|||
|
|
def work(report):
|
|||
|
|
report(done=1, total=2, phase="干活")
|
|||
|
|
return "ok"
|
|||
|
|
|
|||
|
|
task, _ = reg.submit(slot="plug:job", title="推送", func=work)
|
|||
|
|
await asyncio.wait_for(reg.wait(task), 2)
|
|||
|
|
|
|||
|
|
assert not queue.empty()
|
|||
|
|
frames = []
|
|||
|
|
while not queue.empty():
|
|||
|
|
frames.append(queue.get_nowait())
|
|||
|
|
assert frames[0][0]["title"] == "推送"
|
|||
|
|
assert frames[-1][0]["status"] == "done"
|
|||
|
|
# 每帧都是全量快照
|
|||
|
|
assert isinstance(frames[-1], list)
|
|||
|
|
|
|||
|
|
reg.unsubscribe(queue)
|
|||
|
|
reg.unsubscribe(queue) # 幂等
|
|||
|
|
|
|||
|
|
|
|||
|
|
async def test_clear_finished_keeps_running(reg):
|
|||
|
|
gate = threading.Event()
|
|||
|
|
|
|||
|
|
def slow(report):
|
|||
|
|
gate.wait(2)
|
|||
|
|
|
|||
|
|
def quick(report):
|
|||
|
|
return 1
|
|||
|
|
|
|||
|
|
running, _ = reg.submit(slot="plug:slow", title="在跑", func=slow)
|
|||
|
|
done, _ = reg.submit(slot="plug:quick", title="跑完了", func=quick)
|
|||
|
|
await asyncio.wait_for(reg.wait(done), 2)
|
|||
|
|
|
|||
|
|
assert reg.clear_finished() == 1
|
|||
|
|
assert reg.get(done.id) is None
|
|||
|
|
assert reg.get(running.id) is not None
|
|||
|
|
|
|||
|
|
gate.set()
|
|||
|
|
await asyncio.wait_for(reg.wait(running), 2)
|
|||
|
|
|
|||
|
|
|
|||
|
|
async def test_prune_by_max_tasks_keeps_running(reg):
|
|||
|
|
small = type(reg)() # 同类型新实例
|
|||
|
|
small.max_tasks = 2
|
|||
|
|
|
|||
|
|
gate = threading.Event()
|
|||
|
|
|
|||
|
|
def slow(report):
|
|||
|
|
gate.wait(2)
|
|||
|
|
|
|||
|
|
def quick(report):
|
|||
|
|
return 1
|
|||
|
|
|
|||
|
|
running, _ = small.submit(slot="s:run", title="在跑", func=slow)
|
|||
|
|
for i in range(3):
|
|||
|
|
t, _ = small.submit(slot=f"s:{i}", title=f"第{i}次", func=quick)
|
|||
|
|
await asyncio.wait_for(small.wait(t), 2)
|
|||
|
|
|
|||
|
|
assert len(small.list()) <= 2 # 上限生效（终态的被砍）
|
|||
|
|
assert small.get(running.id) is not None # 在跑的绝不动
|
|||
|
|
gate.set()
|
|||
|
|
await asyncio.wait_for(small.wait(running), 2)
|
|||
|
|
|
|||
|
|
|
|||
|
|
async def test_dismiss_only_finished(reg):
|
|||
|
|
"""已结束的能删，在跑的删不掉（删了状态就没人回传了）。"""
|
|||
|
|
gate = threading.Event()
|
|||
|
|
|
|||
|
|
def slow(report):
|
|||
|
|
gate.wait(2)
|
|||
|
|
|
|||
|
|
def quick(report):
|
|||
|
|
return 1
|
|||
|
|
|
|||
|
|
running, _ = reg.submit(slot="plug:slow", title="在跑", func=slow)
|
|||
|
|
done, _ = reg.submit(slot="plug:quick", title="跑完了", func=quick)
|
|||
|
|
await asyncio.wait_for(reg.wait(done), 2)
|
|||
|
|
|
|||
|
|
assert reg.dismiss(running.id) is False # 在跑：拒绝
|
|||
|
|
assert reg.get(running.id) is not None
|
|||
|
|
assert reg.dismiss(done.id) is True
|
|||
|
|
assert reg.get(done.id) is None
|
|||
|
|
assert reg.dismiss(done.id) is False # 已删过 → 再删是 no-op
|
|||
|
|
assert reg.dismiss("不存在的 id") is False
|
|||
|
|
|
|||
|
|
gate.set()
|
|||
|
|
await asyncio.wait_for(reg.wait(running), 2)
|
|||
|
|
|
|||
|
|
|
|||
|
|
async def test_audit_fields_round_trip(reg):
|
|||
|
|
def work(report):
|
|||
|
|
return "ok"
|
|||
|
|
|
|||
|
|
task, _ = reg.submit(
|
|||
|
|
slot="plug:job", title="审计", func=work,
|
|||
|
|
plugin_id="video_analysis", trigger="web", actor="admin", group_id="123456",
|
|||
|
|
)
|
|||
|
|
await asyncio.wait_for(reg.wait(task), 2)
|
|||
|
|
d = task.to_dict()
|
|||
|
|
assert (d["plugin_id"], d["trigger"], d["actor"], d["group_id"]) == (
|
|||
|
|
"video_analysis", "web", "admin", "123456",
|
|||
|
|
)
|
|||
|
|
assert d["status"] == "done" and d["active"] is False
|