chongming-core/test_stream_route.py

"""
网关路由层流式端点单测（依赖隔离：桩掉 chongming 包链、auth 与 chongming_cache）

直接加载 src/chongming/app/core/route_generator.py，用 FastAPI TestClient 验证:
  1. 流式 POST 端点：SSE 输出（media_type / 逐块 data: 帧 / 响应头）
  2. 流式 GET 端点：query 参数正确透传到 payload
  3. 初始错误（无可用 worker）-> HTTP 503
  4. 流中空闲超时 -> SSE error 事件（HTTP 200 已提交）
  5. 兼容性回归：非流式端点输出与改动前一致（JSON 响应）
  6. 非流式端点仍走原 call() 逻辑（不误入流式分支）

运行: uv run python test_stream_route.py
"""
import importlib.util
import logging
import sys
import types
from pathlib import Path

import orjson
from fastapi import FastAPI
from fastapi.testclient import TestClient

REPO_ROOT = Path(__file__).resolve().parent
sys.path.insert(0, str(REPO_ROOT / "src"))
sys.path.insert(0, str(REPO_ROOT / "utils" / "worker" / "src"))

# ---- 依赖隔离：桩掉 chongming 包链 / auth / 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)
_auth_stub = types.ModuleType("chongming.app.core.auth")
_auth_stub.store_user_to_request = lambda x: x  # 仅用于 Depends，桩掉鉴权
_cache_stub = types.ModuleType("chongming_cache")
_cache_stub.Cache = object
for mod in (_pkg, _app, _core, _logger_stub, _auth_stub, _cache_stub):
    sys.modules[mod.__name__] = mod

# 加载真实 worker_manager（route_generator 依赖其 StreamIdleTimeoutError）
_wm_path = REPO_ROOT / "src" / "chongming" / "app" / "core" / "worker_manager.py"
_wm_spec = importlib.util.spec_from_file_location(
    "chongming.app.core.worker_manager", _wm_path
)
_wm_mod = importlib.util.module_from_spec(_wm_spec)
sys.modules["chongming.app.core.worker_manager"] = _wm_mod
_wm_spec.loader.exec_module(_wm_mod)  # type: ignore[union-attr]

# 加载真实 route_generator
_rg_path = REPO_ROOT / "src" / "chongming" / "app" / "core" / "route_generator.py"
_rg_spec = importlib.util.spec_from_file_location(
    "chongming.app.core.route_generator", _rg_path
)
_rg_mod = importlib.util.module_from_spec(_rg_spec)
sys.modules["chongming.app.core.route_generator"] = _rg_mod
_rg_spec.loader.exec_module(_rg_mod)  # type: ignore[union-attr]

register_route_group = _rg_mod.register_route_group
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 FakeWorkerManager:
    """模拟 WorkerManager：call() 同步返回；call_stream() 逐块产出"""

    def __init__(self, chunks=None, stream_error=None):
        self.chunks = chunks if chunks is not None else [{"token": "a"}, {"token": "b"}, {"done": True}]
        self.stream_error = stream_error  # "no_worker" | "idle" | None
        self.call_payloads = []
        self.stream_payloads = []

    async def call(self, worker_type, subject, payload, timeout=3.0, retries=1, cacheable=False):
        self.call_payloads.append((worker_type, subject, payload))
        data = payload["data"]
        return {"result": data.get("a", 0) + data.get("b", 0)}

    async def call_raw(self, worker_type, subject, payload, timeout=3.0, retries=1, cacheable=False):
        # 网关透传路径：payload 为预序列化字节，回复原始字节
        self.call_payloads.append((worker_type, subject, orjson.loads(payload)))
        data = orjson.loads(payload)["data"]
        return orjson.dumps({"result": data.get("a", 0) + data.get("b", 0)})

    async def call_stream(self, worker_type, subject, payload, timeout=300.0, idle_timeout=10.0):
        if self.stream_error == "no_worker":
            raise RuntimeError(f"类型 '{worker_type}' 没有可用 Worker。")
        self.stream_payloads.append((worker_type, subject, payload))

        async def _gen():
            if self.stream_error == "idle":
                yield orjson.dumps({"token": "a"})
                raise StreamIdleTimeoutError("流式空闲超时 (0.1s 无数据)")
            for c in self.chunks:
                yield orjson.dumps(c)

        return _gen()


def build_app(fwm):
    app = FastAPI()
    routes = [
        # 流式 POST（body 参数）
        {
            "subject": "streamtest.chat", "path": "/stream/chat", "method": "POST",
            "params": ["prompt: str", "tokens: int"], "summary": "sse chat",
            "stream": True, "timeout": 300.0, "stream_idle_timeout": 1.0,
            "auth_required": False,
        },
        # 流式 GET（query 参数）
        {
            "subject": "streamtest.progress", "path": "/stream/progress", "method": "GET",
            "params": ["job_id: str", "interval: float"], "summary": "sse progress",
            "stream": True, "timeout": 60.0, "auth_required": False,
        },
        # 非流式对照（兼容性回归）
        {
            "subject": "streamtest.add", "path": "/calc/add", "method": "GET",
            "params": ["a: float", "b: float"], "summary": "add",
            "timeout": 3.0, "auth_required": False,
        },
        {
            "subject": "streamtest.echo", "path": "/calc/echo", "method": "POST",
            "params": ["msg: str"], "summary": "echo",
            "timeout": 3.0, "auth_required": False,
        },
    ]
    register_route_group(app, fwm, "streamtest", routes, prefix="/api/v1")
    return TestClient(app)


def test_stream_post_sse():
    fwm = FakeWorkerManager()
    client = build_app(fwm)
    r = client.post("/api/v1/stream/chat", json={"prompt": "hi", "tokens": 3})
    check("流式 POST 状态码 200", r.status_code == 200, str(r.status_code))
    check(
        "SSE media_type",
        r.headers["content-type"].startswith("text/event-stream"),
        r.headers.get("content-type", ""),
    )
    check(
        "SSE 逐块 data: 帧",
        r.text
        == 'data: {"token":"a"}\n\ndata: {"token":"b"}\n\ndata: {"done":true}\n\n',
        repr(r.text),
    )
    check("防缓冲响应头", r.headers.get("cache-control") == "no-cache")
    check(
        "body 参数透传",
        fwm.stream_payloads[0][2]["data"] == {"prompt": "hi", "tokens": 3},
    )


def test_stream_get_query_params():
    fwm = FakeWorkerManager()
    client = build_app(fwm)
    r = client.get("/api/v1/stream/progress?job_id=job-1&interval=0.5")
    check("流式 GET 状态码 200", r.status_code == 200)
    check(
        "GET query 参数透传",
        fwm.stream_payloads[0][2]["data"] == {"job_id": "job-1", "interval": 0.5},
        str(fwm.stream_payloads[0][2]["data"]),
    )
    check("GET 流式输出 SSE", r.text.startswith("data: "), repr(r.text[:40]))


def test_stream_no_worker_503():
    fwm = FakeWorkerManager(stream_error="no_worker")
    client = build_app(fwm)
    r = client.post("/api/v1/stream/chat", json={"prompt": "hi", "tokens": 1})
    check("初始错误（无 worker）-> 503", r.status_code == 503, str(r.status_code))


def test_stream_idle_timeout_sse_event():
    fwm = FakeWorkerManager(stream_error="idle")
    client = build_app(fwm)
    r = client.post("/api/v1/stream/chat", json={"prompt": "hi", "tokens": 1})
    check("空闲超时保持 200（已提交）", r.status_code == 200)
    check(
        "空闲超时转为 SSE error 事件",
        r.text == 'data: {"token":"a"}\n\ndata: {"error":"流式空闲超时 (0.1s 无数据)"}\n\n',
        repr(r.text),
    )


def test_non_stream_compat():
    """兼容性回归：非流式端点走原 call() 逻辑，输出不变"""
    fwm = FakeWorkerManager()
    client = build_app(fwm)

    r = client.get("/api/v1/calc/add?a=2&b=3")
    check("非流式 GET 仍为 JSON 响应", r.status_code == 200 and r.json() == {"result": 5.0},
          f"status={r.status_code} body={r.text}")

    r = client.post("/api/v1/calc/echo", json={"msg": "hello"})
    check("非流式 POST 仍为 JSON 响应",
          r.status_code == 200 and r.json() == {"result": 0.0},
          f"status={r.status_code} body={r.text}")

    check("非流式调用未进入 call_stream", len(fwm.stream_payloads) == 0)
    check("非流式调用仍走 call_raw", len(fwm.call_payloads) == 2)


def run():
    test_stream_post_sse()
    test_stream_get_query_params()
    test_stream_no_worker_503()
    test_stream_idle_timeout_sse_event()
    test_non_stream_compat()

    print(f"\n结果: {len(failures)} 个失败")
    return 1 if failures else 0


if __name__ == "__main__":
    sys.exit(run())