119 lines
3.2 KiB
Python
119 lines
3.2 KiB
Python
|
"""
|
|||
|
|
Worker 框架流式回复单测（不依赖真实 zenoh 会话）
|
|||
|
|
|
|||
|
|
直接测 runtime.stream_replies_to_query（fake query 仅需 key_expr 与 reply）:
|
|||
|
|
1. 异步生成器：多块顺序 reply，payload 为 orjson JSON
|
|||
|
|
2. 同步生成器：同样逐块 reply
|
|||
|
|
3. 生成器中途异常：{"error": ...} 收尾块后停流
|
|||
|
|
4. reply 失败（查询已终止）：立即停流且不抛异常
|
|||
|
|
|
|||
|
|
运行: uv run python test_stream_worker.py
|
|||
|
|
"""
|
|||
|
|
import asyncio
|
|||
|
|
import sys
|
|||
|
|
from pathlib import Path
|
|||
|
|
|
|||
|
|
import orjson
|
|||
|
|
|
|||
|
|
REPO_ROOT = Path(__file__).resolve().parent
|
|||
|
|
# 优先使用本地框架 checkout（即使 chongming_worker 已通过 uv 安装也以本地为准）
|
|||
|
|
sys.path.insert(0, str(REPO_ROOT / "utils" / "worker" / "src"))
|
|||
|
|
|
|||
|
|
from chongming_worker.runtime import stream_replies_to_query # noqa: E402
|
|||
|
|
|
|||
|
|
failures = []
|
|||
|
|
|
|||
|
|
|
|||
|
|
def check(name, cond, detail=""):
|
|||
|
|
status = "PASS" if cond else "FAIL"
|
|||
|
|
print(f"[{status}] {name}" + (f" — {detail}" if detail else ""))
|
|||
|
|
if not cond:
|
|||
|
|
failures.append(name)
|
|||
|
|
|
|||
|
|
|
|||
|
|
class FakeQuery:
|
|||
|
|
"""仅实现 stream_replies_to_query 用到的两个接口（key_expr / reply）"""
|
|||
|
|
|
|||
|
|
def __init__(self, key_expr="worker/example/worker-0/example.stream_chat"):
|
|||
|
|
self.key_expr = key_expr
|
|||
|
|
self.replies = []
|
|||
|
|
self.fail_reply = False
|
|||
|
|
|
|||
|
|
def reply(self, key_expr, payload):
|
|||
|
|
if self.fail_reply:
|
|||
|
|
raise RuntimeError("query terminated")
|
|||
|
|
self.replies.append((key_expr, payload))
|
|||
|
|
|
|||
|
|
|
|||
|
|
def test_async_generator():
|
|||
|
|
q = FakeQuery()
|
|||
|
|
|
|||
|
|
async def gen():
|
|||
|
|
yield {"token": "a"}
|
|||
|
|
yield {"token": "b"}
|
|||
|
|
|
|||
|
|
asyncio.run(stream_replies_to_query(q, gen()))
|
|||
|
|
check("异步生成器逐块回复", len(q.replies) == 2)
|
|||
|
|
check("回复 key_expr 正确", all(k == q.key_expr for k, _ in q.replies))
|
|||
|
|
check(
|
|||
|
|
"payload 为 JSON dict",
|
|||
|
|
[orjson.loads(p) for _, p in q.replies] == [{"token": "a"}, {"token": "b"}],
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
|
|||
|
|
def test_sync_generator():
|
|||
|
|
q = FakeQuery()
|
|||
|
|
|
|||
|
|
def gen():
|
|||
|
|
yield {"n": 1}
|
|||
|
|
yield {"n": 2}
|
|||
|
|
|
|||
|
|
asyncio.run(stream_replies_to_query(q, gen()))
|
|||
|
|
check(
|
|||
|
|
"同步生成器逐块回复",
|
|||
|
|
[orjson.loads(p) for _, p in q.replies] == [{"n": 1}, {"n": 2}],
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
|
|||
|
|
def test_mid_stream_exception():
|
|||
|
|
q = FakeQuery()
|
|||
|
|
|
|||
|
|
async def gen():
|
|||
|
|
yield {"token": "a"}
|
|||
|
|
raise ValueError("boom")
|
|||
|
|
|
|||
|
|
asyncio.run(stream_replies_to_query(q, gen()))
|
|||
|
|
check("中途异常收尾 error 块", len(q.replies) == 2)
|
|||
|
|
check("error 块内容", orjson.loads(q.replies[1][1]) == {"error": "boom"})
|
|||
|
|
|
|||
|
|
|
|||
|
|
def test_reply_failure_stops():
|
|||
|
|
q = FakeQuery()
|
|||
|
|
q.fail_reply = True
|
|||
|
|
yielded = []
|
|||
|
|
|
|||
|
|
async def gen():
|
|||
|
|
yielded.append("yield-1")
|
|||
|
|
yield {"token": "a"}
|
|||
|
|
yielded.append("yield-2") # 首个块 reply 失败后不应被执行到
|
|||
|
|
yield {"token": "b"}
|
|||
|
|
|
|||
|
|
# 不应抛异常，且首个块 reply 失败后立即停流
|
|||
|
|
asyncio.run(stream_replies_to_query(q, gen()))
|
|||
|
|
check("reply 失败不抛异常", True)
|
|||
|
|
check("reply 失败后立即停流", yielded == ["yield-1"])
|
|||
|
|
|
|||
|
|
|
|||
|
|
def run():
|
|||
|
|
test_async_generator()
|
|||
|
|
test_sync_generator()
|
|||
|
|
test_mid_stream_exception()
|
|||
|
|
test_reply_failure_stops()
|
|||
|
|
|
|||
|
|
print(f"\n结果: {len(failures)} 个失败")
|
|||
|
|
return 1 if failures else 0
|
|||
|
|
|
|||
|
|
|
|||
|
|
if __name__ == "__main__":
|
|||
|
|
sys.exit(run())
|