241 lines
9.3 KiB
Python
241 lines
9.3 KiB
Python
"""
|
||
网关路由层流式端点单测（依赖隔离：桩掉 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.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
|
||
_cache_stub.DistributedLock = object # worker_manager 顶层导入（本测试不触达）
|
||
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]
|
||
|
||
# 加载真实 call_stats（route_generator 依赖其 spawn/record_call）
|
||
_cs_path = REPO_ROOT / "src" / "chongming" / "app" / "core" / "call_stats.py"
|
||
_cs_spec = importlib.util.spec_from_file_location(
|
||
"chongming.app.core.call_stats", _cs_path
|
||
)
|
||
_cs_mod = importlib.util.module_from_spec(_cs_spec)
|
||
sys.modules["chongming.app.core.call_stats"] = _cs_mod
|
||
_cs_spec.loader.exec_module(_cs_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,
|
||
instance_id=None):
|
||
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,
|
||
instance_id=None):
|
||
# 网关透传路径：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,
|
||
instance_id=None):
|
||
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())
|