modify: 测试
This commit is contained in:
@@ -0,0 +1,180 @@
|
||||
"""
|
||||
真实 Zenoh 会话下的流式端到端测试（peer 模式，无需 router）
|
||||
|
||||
单进程内起两个 Zenoh 会话（listen + connect），跑通完整生产链路:
|
||||
@rpc_method(stream=True) 注册 -> create_worker 声明 Queryable
|
||||
-> WorkerManager.call_stream 消费多回复 -> 逐块到达
|
||||
|
||||
验证:
|
||||
1. 流式方法：3+1 块按序到达（真实 Zenoh 多回复 + ConsolidationMode.NONE）
|
||||
2. 流式中途异常：{"error": ...} 收尾块转发后终止
|
||||
3. 非流式方法：仍是单回复（兼容性回归）
|
||||
4. generate_route_configs 输出 stream 标记、非流式不含
|
||||
5. decorators 对生成器函数的调用不消费生成器
|
||||
|
||||
运行: uv run python test_stream_e2e_zenoh.py
|
||||
"""
|
||||
import asyncio
|
||||
import importlib.util
|
||||
import logging
|
||||
import sys
|
||||
import types
|
||||
from pathlib import Path
|
||||
|
||||
import orjson
|
||||
import zenoh
|
||||
|
||||
REPO_ROOT = Path(__file__).resolve().parent
|
||||
sys.path.insert(0, str(REPO_ROOT / "utils" / "worker" / "src"))
|
||||
sys.path.insert(0, str(REPO_ROOT / "src"))
|
||||
|
||||
# ---- 网关依赖隔离（同 test_stream_gateway）：桩掉 chongming 包链与 chongming_cache ----
|
||||
_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
|
||||
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
|
||||
|
||||
from chongming_worker import clear_registry, rpc_method # noqa: E402
|
||||
from chongming_worker.decorators import generate_route_configs # noqa: E402
|
||||
from chongming_worker.runtime import create_worker # 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)
|
||||
|
||||
|
||||
LISTEN_PORT = 7447
|
||||
LISTEN = f"tcp/127.0.0.1:{LISTEN_PORT}"
|
||||
|
||||
|
||||
def open_session(listen=None):
|
||||
"""peer 模式会话：worker 侧 listen，网关侧 connect（显式端点，禁用组播）"""
|
||||
c = zenoh.Config()
|
||||
c.insert_json5("scouting/multicast/enabled", "false")
|
||||
if listen:
|
||||
c.insert_json5("listen/endpoints", f'["{listen}"]')
|
||||
else:
|
||||
c.insert_json5("connect/endpoints", f'["{LISTEN}"]')
|
||||
return zenoh.open(c)
|
||||
|
||||
|
||||
async def run_checks():
|
||||
clear_registry()
|
||||
|
||||
# ---- 流式端点（返回异步生成器）----
|
||||
@rpc_method(subject="streamtest.echo", method="POST", stream=True)
|
||||
async def echo_stream(prompt: str = "hi", tokens: int = 3):
|
||||
for i in range(tokens):
|
||||
await asyncio.sleep(0.01)
|
||||
yield {"token": f"{prompt}-{i}"}
|
||||
yield {"done": True}
|
||||
|
||||
# ---- 流式端点：中途异常（框架应发 {"error": ...} 收尾块）----
|
||||
@rpc_method(subject="streamtest.fail", method="POST", stream=True)
|
||||
async def fail_stream():
|
||||
yield {"token": "a"}
|
||||
raise ValueError("mid-boom")
|
||||
|
||||
# ---- 非流式端点（兼容性对照）----
|
||||
@rpc_method(subject="streamtest.add", method="GET")
|
||||
async def add(a: float, b: float):
|
||||
return a + b
|
||||
|
||||
# 1. 配置生成：流式标记
|
||||
configs = generate_route_configs("streamtest")
|
||||
stream_cfg = next(c for c in configs if c["subject"] == "streamtest.echo")
|
||||
check("generate_route_configs 输出 stream 标记", stream_cfg.get("stream") is True)
|
||||
add_cfg = next(c for c in configs if c["subject"] == "streamtest.add")
|
||||
check("非流式配置不含 stream 标记", "stream" not in add_cfg)
|
||||
|
||||
# 2. 真实 Zenoh 会话（peer 模式，无 router）
|
||||
session_worker = open_session(listen=LISTEN)
|
||||
session_gw = open_session()
|
||||
|
||||
# 3. 启动真实 Worker（create_worker 声明 Queryable）
|
||||
worker_task = asyncio.create_task(
|
||||
create_worker(session_worker, "streamtest", "w1")
|
||||
)
|
||||
await asyncio.sleep(0.5) # 等待 queryable 在 peer 间可见
|
||||
|
||||
try:
|
||||
# 4. 网关 WorkerManager.call_stream 真实消费流式端点
|
||||
wm = WorkerManager(session_gw)
|
||||
wm.register_worker("streamtest", "w1")
|
||||
|
||||
async def _consume(subject, data, idle_timeout=2.0):
|
||||
gen = await wm.call_stream(
|
||||
"streamtest", subject,
|
||||
{"meta": {"method": "POST"}, "data": data},
|
||||
timeout=10.0, idle_timeout=idle_timeout,
|
||||
)
|
||||
return [orjson.loads(b) async for b in gen]
|
||||
|
||||
chunks = await _consume("streamtest.echo", {"prompt": "hi", "tokens": 3})
|
||||
check(
|
||||
"流式 3+1 块按序到达（真实多回复）",
|
||||
chunks == [
|
||||
{"token": "hi-0"}, {"token": "hi-1"}, {"token": "hi-2"}, {"done": True},
|
||||
],
|
||||
f"chunks={chunks}",
|
||||
)
|
||||
|
||||
# 5. 流式中途异常：错误块收尾
|
||||
chunks = await _consume("streamtest.fail", {})
|
||||
check(
|
||||
"流式中途异常 error 块收尾",
|
||||
chunks == [{"token": "a"}, {"error": "mid-boom"}],
|
||||
f"chunks={chunks}",
|
||||
)
|
||||
|
||||
# 6. 非流式兼容：call() 仍单回复
|
||||
result = await wm.call(
|
||||
"streamtest", "streamtest.add",
|
||||
{"meta": {"method": "GET"}, "data": {"a": 2, "b": 3}},
|
||||
timeout=10.0,
|
||||
)
|
||||
check("非流式 call() 正常单回复", result == 5.0, f"result={result}")
|
||||
|
||||
# 7. 负载均衡路径兼容：liveliness 发现的 worker 也能流式（同 call）
|
||||
worker_expr = await wm.get_worker("streamtest")
|
||||
check("get_worker 仍可解析本地 worker", worker_expr is not None)
|
||||
|
||||
finally:
|
||||
worker_task.cancel()
|
||||
await asyncio.gather(worker_task, return_exceptions=True)
|
||||
session_gw.close()
|
||||
session_worker.close()
|
||||
|
||||
|
||||
def run():
|
||||
asyncio.run(run_checks())
|
||||
print(f"\n结果: {len(failures)} 个失败")
|
||||
return 1 if failures else 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(run())
|
||||
Reference in New Issue
Block a user