modify: 流式输出

test_stream_worker.py

@@ -0,0 +1,118 @@
+"""
+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())