chongming-core/test/test_stream_e2e_zenoh.py

"""
真实 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.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
_cache_stub.DistributedLock = object  # worker_manager 顶层导入（本测试不触达）
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)


import os
LISTEN_PORT = int(os.environ.get("CM_STREAM_TEST_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())