""" 网关路由层流式端点单测（依赖隔离：桩掉 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())