248 lines
7.6 KiB
Python
248 lines
7.6 KiB
Python
"""
|
||
网关 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())
|