chongming-core/test_stream_gateway.py

"""
网关 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())