208 lines
9.1 KiB
Python
208 lines
9.1 KiB
Python
"""
|
||
HTTP 实例定向头（X-Chongming-Instance）透传测试
|
||
|
||
依赖隔离：桩掉 chongming 包链、auth 与 chongming_cache，直接加载真实
|
||
route_generator.py 与 worker_manager.py，用 FastAPI TestClient 验证:
|
||
1. 携带 X-Chongming-Instance 的请求，该实例 id 透传到 call_raw / call_stream
|
||
2. 未携带该头时 instance_id 为 None（行为与改动前完全一致，兼容旧调用方）
|
||
3. 覆盖全部端点形态：GET query / POST body / validate=false 原样透传 / SSE 流式
|
||
4. 头值为空白时视为未指定
|
||
5. 定向到离线实例（RuntimeError）-> HTTP 503，错误信息透出
|
||
|
||
运行: .venv/bin/python test/test_instance_targeting_http.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
|
||
|
||
TEST_DIR = Path(__file__).resolve().parent
|
||
REPO_ROOT = TEST_DIR.parent
|
||
sys.path.insert(0, str(REPO_ROOT / "gateway" / "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
|
||
_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
|
||
|
||
_wm_path = REPO_ROOT / "gateway" / "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]
|
||
|
||
_cs_path = REPO_ROOT / "gateway" / "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]
|
||
|
||
_rg_path = REPO_ROOT / "gateway" / "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
|
||
INSTANCE_HEADER = _rg_mod.INSTANCE_HEADER
|
||
|
||
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:
|
||
"""记录每次调用收到的 instance_id；可指定定点失败以验证 503 映射。"""
|
||
|
||
def __init__(self, fail_instance=None):
|
||
self.fail_instance = fail_instance # 指定实例报错（模拟实例离线）
|
||
self.raw_instances = [] # call_raw 收到的 instance_id
|
||
self.stream_instances = [] # call_stream 收到的 instance_id
|
||
|
||
async def call_raw(self, worker_type, subject, payload, timeout=3.0,
|
||
retries=1, cacheable=False, instance_id=None):
|
||
self.raw_instances.append(instance_id)
|
||
if instance_id and instance_id == self.fail_instance:
|
||
raise RuntimeError(
|
||
f"类型 '{worker_type}' 的实例 '{instance_id}' 没有在线副本"
|
||
"（定向调用不跨实例回退）。"
|
||
)
|
||
return orjson.dumps({"op": subject, "instance": instance_id})
|
||
|
||
async def call_stream(self, worker_type, subject, payload, timeout=300.0,
|
||
idle_timeout=10.0, instance_id=None):
|
||
self.stream_instances.append(instance_id)
|
||
if instance_id and instance_id == self.fail_instance:
|
||
raise RuntimeError(f"实例 '{instance_id}' 没有在线副本")
|
||
|
||
async def _gen():
|
||
yield orjson.dumps({"instance": instance_id})
|
||
|
||
return _gen()
|
||
|
||
|
||
ROUTES = [
|
||
{"subject": "t.get", "path": "/calc/add", "method": "GET",
|
||
"params": ["a: float", "b: float"], "summary": "add",
|
||
"timeout": 3.0, "auth_required": False},
|
||
{"subject": "t.post", "path": "/calc/echo", "method": "POST",
|
||
"params": ["msg: str"], "summary": "echo",
|
||
"timeout": 3.0, "auth_required": False},
|
||
{"subject": "t.raw", "path": "/calc/raw", "method": "POST",
|
||
"params": ["payload: any"], "summary": "原样透传",
|
||
"validate": False, "timeout": 3.0, "auth_required": False},
|
||
{"subject": "t.stream", "path": "/stream/chat", "method": "POST",
|
||
"params": ["prompt: str"], "summary": "sse",
|
||
"stream": True, "timeout": 60.0, "stream_idle_timeout": 1.0,
|
||
"auth_required": False},
|
||
]
|
||
|
||
|
||
def build_client(fwm):
|
||
app = FastAPI()
|
||
register_route_group(app, fwm, "demotype", ROUTES, prefix="/api/v1")
|
||
return TestClient(app)
|
||
|
||
|
||
def main():
|
||
# ── 1. GET query 端点 ────────────────────────────────────
|
||
fwm = FakeWorkerManager()
|
||
c = build_client(fwm)
|
||
r = c.get("/api/v1/calc/add?a=1&b=2", headers={INSTANCE_HEADER: "edge-01"})
|
||
check("1.1 GET 携带头 -> 200", r.status_code == 200, f"status={r.status_code}")
|
||
check("1.2 GET 头透传到 call_raw", fwm.raw_instances == ["edge-01"],
|
||
f"instances={fwm.raw_instances}")
|
||
check("1.3 响应体回带实例（证明打到目标实例）",
|
||
r.json().get("instance") == "edge-01", f"body={r.json()}")
|
||
|
||
# ── 2. 未携带头：兼容旧调用方 ────────────────────────────
|
||
fwm = FakeWorkerManager()
|
||
c = build_client(fwm)
|
||
r = c.get("/api/v1/calc/add?a=1&b=2")
|
||
check("2.1 未携带头 -> 200（默认随机均衡路径不变）", r.status_code == 200)
|
||
check("2.2 未携带头时 instance_id=None", fwm.raw_instances == [None],
|
||
f"instances={fwm.raw_instances}")
|
||
check("2.3 响应体不含实例", r.json().get("instance") is None, f"body={r.json()}")
|
||
|
||
# ── 3. POST body 端点 ────────────────────────────────────
|
||
fwm = FakeWorkerManager()
|
||
c = build_client(fwm)
|
||
r = c.post("/api/v1/calc/echo", json={"msg": "hi"},
|
||
headers={INSTANCE_HEADER: "edge-02"})
|
||
check("3.1 POST body 头透传", fwm.raw_instances == ["edge-02"],
|
||
f"instances={fwm.raw_instances}")
|
||
|
||
# ── 4. validate=false 原样透传端点 ───────────────────────
|
||
fwm = FakeWorkerManager()
|
||
c = build_client(fwm)
|
||
r = c.post("/api/v1/calc/raw", content=b'{"anything": [1, 2]}',
|
||
headers={INSTANCE_HEADER: "edge-03",
|
||
"content-type": "application/json"})
|
||
check("4.1 原样透传端点头透传", fwm.raw_instances == ["edge-03"],
|
||
f"instances={fwm.raw_instances}")
|
||
|
||
# ── 5. SSE 流式端点 ──────────────────────────────────────
|
||
fwm = FakeWorkerManager()
|
||
c = build_client(fwm)
|
||
with c.stream("POST", "/api/v1/stream/chat", json={"prompt": "x"},
|
||
headers={INSTANCE_HEADER: "edge-04"}) as r:
|
||
body = b"".join(r.iter_bytes())
|
||
check("5.1 流式端点头透传", fwm.stream_instances == ["edge-04"],
|
||
f"instances={fwm.stream_instances}")
|
||
check("5.2 流式响应含目标实例", b"edge-04" in body, f"body={body[:80]!r}")
|
||
|
||
# ── 6. 空白头值视为未指定 ────────────────────────────────
|
||
fwm = FakeWorkerManager()
|
||
c = build_client(fwm)
|
||
r = c.get("/api/v1/calc/add?a=1&b=2", headers={INSTANCE_HEADER: " "})
|
||
check("6.1 空白头值 -> instance_id=None", fwm.raw_instances == [None],
|
||
f"instances={fwm.raw_instances}")
|
||
|
||
# ── 7. 定向实例离线 -> 503 ───────────────────────────────
|
||
fwm = FakeWorkerManager(fail_instance="edge-gone")
|
||
c = build_client(fwm)
|
||
r = c.get("/api/v1/calc/add?a=1&b=2", headers={INSTANCE_HEADER: "edge-gone"})
|
||
check("7.1 定向实例离线 -> 503", r.status_code == 503, f"status={r.status_code}")
|
||
check("7.2 503 详情点名实例", "edge-gone" in str(r.json().get("detail", "")),
|
||
f"detail={r.json().get('detail')}")
|
||
|
||
print()
|
||
if failures:
|
||
print(f"FAILED: {len(failures)} 项 -> {failures}")
|
||
return 1
|
||
print("全部通过")
|
||
return 0
|
||
|
||
|
||
if __name__ == "__main__":
|
||
raise SystemExit(main())
|