""" 网关 call_stream 单测（依赖隔离：桩掉 chongming 包链与 chongming_cache） 直接加载 gateway/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.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 仅作类型标注使用 _cache_stub.DistributedLock = object # 同上（set_if_not_exists 用的分布式锁） for mod in (_pkg, _app, _core, _logger_stub, _cache_stub): sys.modules[mod.__name__] = mod _wm_path = REPO_ROOT / "gateway" / "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())