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())
|