modify: 流式输出
This commit is contained in:
@@ -0,0 +1,221 @@
|
||||
"""
|
||||
网关路由层流式端点单测（依赖隔离：桩掉 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_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()", 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())
|
||||
Reference in New Issue
Block a user