modify: 测试

2026-08-23 05:10:27 +08:00
parent a1be33376e
commit eeffc06841
24 changed files with 2 additions and 2 deletions
@@ -0,0 +1,246 @@
+"""
+网关 call_stream 单测（依赖隔离：桩掉 chongming 包链与 chongming_cache）
+
+直接加载 src/chongming/app/core/worker_manager.py 模块（避开 chongming.app
+整链导入与 config.toml 依赖），验证:
+  1. 多块顺序产出（ConsolidationMode.NONE 逐块消费）
+  2. Worker 错误收尾块 {"error": ...}：转发后终止流
+  3. zenoh 错误回复（reply.err）：转发错误块后终止流
+  4. 空闲超时：抛 StreamIdleTimeoutError
+  5. 无可用 worker：抛 RuntimeError
+  6. 兼容性回归：call()（非流式路径）仍只取第一个回复
+
+运行: uv run python test_stream_gateway.py
+"""
+import asyncio
+import importlib.util
+import logging
+import sys
+import time
+import types
+from pathlib import Path
+from types import SimpleNamespace
+
+import orjson
+import zenoh
+
+REPO_ROOT = Path(__file__).resolve().parent
+
+# ---- 依赖隔离：桩掉网关包链（chongming.app 顶层导入 fastapi/uvicorn/config）----
+_pkg = types.ModuleType("chongming")
+_pkg.__path__ = []
+_app = types.ModuleType("chongming.app")
+_app.__path__ = []
+_core = types.ModuleType("chongming.app.core")
+_core.__path__ = []
+_logger_stub = types.ModuleType("chongming.app.core.logger")
+_logger_stub.get_logger = lambda name: logging.getLogger(name)
+_cache_stub = types.ModuleType("chongming_cache")
+_cache_stub.Cache = object  # WorkerManager 仅作类型标注使用
+for mod in (_pkg, _app, _core, _logger_stub, _cache_stub):
+    sys.modules[mod.__name__] = mod
+
+_wm_path = REPO_ROOT / "src" / "chongming" / "app" / "core" / "worker_manager.py"
+_spec = importlib.util.spec_from_file_location(
+    "chongming.app.core.worker_manager", _wm_path
+)
+_wm_mod = importlib.util.module_from_spec(_spec)
+sys.modules["chongming.app.core.worker_manager"] = _wm_mod
+_spec.loader.exec_module(_wm_mod)  # type: ignore[union-attr]
+
+WorkerManager = _wm_mod.WorkerManager
+StreamIdleTimeoutError = _wm_mod.StreamIdleTimeoutError
+
+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 FakeReply:
+    """模拟 zenoh.Reply：ok / err 二选一"""
+
+    def __init__(self, payload=None, error=None):
+        if error is None:
+            self.ok = SimpleNamespace(payload=payload)
+            self.err = None
+        else:
+            self.ok = None
+            self.err = SimpleNamespace(payload=error)
+
+
+class FakeSession:
+    """模拟 zenoh.Session.get()：返回 reply 迭代器（可迭代对象）"""
+
+    def __init__(self, reply_source):
+        self._reply_source = reply_source
+        self.get_calls = []
+
+    def get(self, key_expr, payload, consolidation=None, timeout=None):
+        self.get_calls.append((str(key_expr), payload, consolidation, timeout))
+        return iter(self._reply_source)
+
+
+def make_manager(replies, worker_id="worker-0"):
+    wm = WorkerManager(FakeSession(replies))
+    wm.register_worker("example", worker_id)
+    return wm
+
+
+def test_idle_timeout():
+    # 模拟阻塞式 reply 迭代器：每个 __next__ 阻塞 0.3s（首个块 0.3s 后才到达），
+    # 而空闲超时仅 0.05s -> 应抛 StreamIdleTimeoutError
+    class BlockingReplyIterator:
+        def __init__(self):
+            self._done = False
+
+        def __iter__(self):
+            return self
+
+        def __next__(self):
+            time.sleep(0.3)
+            if self._done:
+                raise StopIteration
+            self._done = True
+            return FakeReply(payload=orjson.dumps({"token": "slow"}))
+
+    wm = WorkerManager(FakeSession(BlockingReplyIterator()))
+
+
+def collect_stream(wm, idle_timeout=1.0):
+    """消费 call_stream 全部块并返回解码后的列表"""
+
+    async def _run():
+        gen = await wm.call_stream(
+            "example", "stream_chat", {"meta": {"method": "POST"}, "data": {}},
+            timeout=5.0, idle_timeout=idle_timeout,
+        )
+        return [orjson.loads(b) async for b in gen]
+
+    return asyncio.run(_run())
+
+
+def test_multiple_chunks():
+    replies = [
+        FakeReply(payload=orjson.dumps({"token": "a"})),
+        FakeReply(payload=orjson.dumps({"token": "b"})),
+    ]
+    wm = make_manager(replies)
+    chunks = collect_stream(wm)
+    check("多块顺序产出", chunks == [{"token": "a"}, {"token": "b"}])
+    check(
+        "使用 ConsolidationMode.NONE",
+        wm.session.get_calls[0][2] is zenoh.ConsolidationMode.NONE,
+    )
+    check("key_expr 正确", wm.session.get_calls[0][0] == "worker/example/worker-0/stream_chat")
+
+
+def test_error_chunk_terminates():
+    replies = [
+        FakeReply(payload=orjson.dumps({"token": "a"})),
+        FakeReply(payload=orjson.dumps({"error": "boom"})),
+        FakeReply(payload=orjson.dumps({"token": "c"})),  # 不应被消费
+    ]
+    wm = make_manager(replies)
+    chunks = collect_stream(wm)
+    check("错误块转发后终止", chunks == [{"token": "a"}, {"error": "boom"}])
+
+
+def test_err_reply_terminates():
+    replies = [
+        FakeReply(payload=orjson.dumps({"token": "a"})),
+        FakeReply(error="oops"),
+    ]
+    wm = make_manager(replies)
+    chunks = collect_stream(wm)
+    check(
+        "reply.err 转发错误块后终止",
+        chunks == [{"token": "a"}, {"error": "Worker worker/example/worker-0 返回错误: oops"}],
+    )
+
+
+def test_idle_timeout():
+    # 模拟阻塞式 reply 迭代器：每个 __next__ 阻塞 0.3s（首个块 0.3s 后才到达），
+    # 而空闲超时仅 0.05s -> 应抛 StreamIdleTimeoutError
+    class BlockingReplyIterator:
+        def __init__(self):
+            self._done = False
+
+        def __iter__(self):
+            return self
+
+        def __next__(self):
+            time.sleep(0.3)
+            if self._done:
+                raise StopIteration
+            self._done = True
+            return FakeReply(payload=orjson.dumps({"token": "slow"}))
+
+    wm = WorkerManager(FakeSession(BlockingReplyIterator()))
+    wm.register_worker("example", "worker-0")
+
+    async def _run():
+        gen = await wm.call_stream(
+            "example", "stream_chat", {}, timeout=5.0, idle_timeout=0.05
+        )
+        async for _ in gen:
+            pass
+
+    try:
+        asyncio.run(_run())
+        check("空闲超时抛 StreamIdleTimeoutError", False)
+    except StreamIdleTimeoutError as e:
+        check("空闲超时抛 StreamIdleTimeoutError", True, str(e))
+
+
+def test_no_worker():
+    wm = WorkerManager(FakeSession([]))  # 未注册任何 worker
+
+    async def _run():
+        gen = await wm.call_stream("missing", "x", {})
+        async for _ in gen:
+            pass
+
+    try:
+        asyncio.run(_run())
+        check("无 worker 抛 RuntimeError", False)
+    except RuntimeError as e:
+        check("无 worker 抛 RuntimeError", True, str(e)[:30])
+
+
+def test_call_compat_first_reply_only():
+    """兼容性回归：非流式 call() 仍只取第一个回复，行为不变"""
+    replies = [
+        FakeReply(payload=orjson.dumps({"result": 1})),
+        FakeReply(payload=orjson.dumps({"result": 2})),
+    ]
+    wm = make_manager(replies)
+
+    async def _run():
+        return await wm.call(
+            "example", "stream_chat", {"meta": {"method": "POST"}, "data": {}},
+            timeout=5.0,
+        )
+
+    result = asyncio.run(_run())
+    check("call() 仍只取第一个回复（非流式兼容）", result == {"result": 1})
+
+
+def run():
+    test_multiple_chunks()
+    test_error_chunk_terminates()
+    test_err_reply_terminates()
+    test_idle_timeout()
+    test_no_worker()
+    test_call_compat_first_reply_only()
+
+    print(f"\n结果: {len(failures)} 个失败")
+    return 1 if failures else 0
+
+
+if __name__ == "__main__":
+    sys.exit(run())