""" 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 / "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 / "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 / "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 / "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())