"""Stub Worker — 模拟高性能后端，用于网关开发与测试。 每个 stub 监听一个端口，提供： - GET /health → 200（可用 ?fail=1 模拟故障） - POST /api/v1/face/measure → mock 响应（含 annotated_image_base64） - POST /api/v1/hair/grow → mock 响应 - POST /api/v1/hair/grow-b → mock 响应 - POST /api/v1/face/features → mock 响应 - POST /api/v1/hairline/generate → mock 响应 鉴权：校验 X-Internal-Token 头，不匹配返回 401。 并发模拟：可通过 ?delay=N 让接口 sleep N 秒（默认 1）。 故障模拟：GET /health?fail=1 返回 503。 用法： python gateway/stub_worker.py --port 28187 python gateway/stub_worker.py --port 28188 """ import argparse import base64 import json import os import sys import time import uuid from typing import Optional # 确保项目根在 sys.path sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) import uvicorn from fastapi import FastAPI, File, Form, Header, Request, UploadFile from fastapi.responses import JSONResponse # --------------------------------------------------------------------------- # 一小张 PNG 用于 mock base64 返回 # 1x1 白色像素 PNG（最小合法 PNG） # --------------------------------------------------------------------------- TINY_PNG_BASE64 = ( "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8/5+hHgAHggJ/PchI7wAAAABJRU5ErkJggg==" ) TINY_PNG_BYTES = base64.b64decode(TINY_PNG_BASE64) # --------------------------------------------------------------------------- # 应用工厂 # --------------------------------------------------------------------------- def create_app(port: int, password: str) -> FastAPI: app = FastAPI( title=f"Stub Worker :{port}", version="0.1.0", docs_url=None, redoc_url=None, ) # --- 鉴权依赖 --- def check_token(x_internal_token: Optional[str] = Header(default=None)): if x_internal_token != password: return False return True # --- 通用 mock data 生成 --- def mock_ok(data: dict, delay: float = 1.0): """模拟处理延迟后返回标准响应。""" if delay > 0: time.sleep(delay) return { "code": 0, "message": "success", "request_id": f"stub-{port}-{uuid.uuid4().hex[:8]}", "data": data, } def mock_err(code: int, message: str): return { "code": code, "message": message, "request_id": f"stub-{port}-{uuid.uuid4().hex[:8]}", "data": None, } # --- /health --- @app.get("/health", include_in_schema=False) async def health(request: Request, x_internal_token: Optional[str] = Header(default=None)): # /health 也校验 token（生产环境 worker 行为） if x_internal_token != password: return JSONResponse(status_code=401, content={"detail": "unauthorized"}) fail = request.query_params.get("fail") if fail == "1": return JSONResponse(status_code=503, content={"status": "error"}) return {"status": "ok", "worker": f"stub-{port}"} # --- 接口1 --- @app.post("/api/v1/face/measure") async def face_measure( request: Request, x_internal_token: Optional[str] = Header(default=None), image_file: Optional[UploadFile] = File(default=None), image_url: Optional[str] = Form(default=None), image_base64: Optional[str] = Form(default=None), ): if x_internal_token != password: return JSONResponse(status_code=401, content={"detail": "unauthorized"}) delay = float(request.query_params.get("delay", "1")) data = { "annotated_image_base64": f"data:image/png;base64,{TINY_PNG_BASE64}", "face_total_height_cm": 13.76, "four_courts": { "top_court_cm": 3.44, "upper_court_cm": 3.44, "middle_court_cm": 3.44, "lower_court_cm": 3.44, "ratios": { "top_court": 0.25, "upper_court": 0.25, "middle_court": 0.25, "lower_court": 0.25, }, }, "seven_eyes": { "eye_width_cm": 3.44, "face_width_cm": 24.08, "inter_eye_distance_cm": 3.44, "ratios": {"eye_width": 0.143, "inter_eye_distance": 0.143}, }, "landmarks": { "hair_top": {"x": 540, "y": 120}, "hairline": {"x": 540, "y": 430}, "brow_center": {"x": 540, "y": 740}, "nose_bottom": {"x": 540, "y": 1050}, "chin_tip": {"x": 540, "y": 1360}, }, } return mock_ok(data, delay=delay) # --- 接口2 --- @app.post("/api/v1/hair/grow") async def hair_grow( request: Request, x_internal_token: Optional[str] = Header(default=None), image_file: Optional[UploadFile] = File(default=None), image_url: Optional[str] = Form(default=None), image_base64: Optional[str] = Form(default=None), beauty_enabled: bool = Form(default=False), ): if x_internal_token != password: return JSONResponse(status_code=401, content={"detail": "unauthorized"}) delay = float(request.query_params.get("delay", "1")) data = { "results": [ { "image_base64": f"data:image/png;base64,{TINY_PNG_BASE64}", "hairline_type": "花瓣形", "order": 1, }, { "image_base64": f"data:image/png;base64,{TINY_PNG_BASE64}", "hairline_type": "波浪形", "order": 2, }, ] } return mock_ok(data, delay=delay) # --- 接口3 --- @app.post("/api/v1/hair/grow-b") async def hair_grow_b( request: Request, x_internal_token: Optional[str] = Header(default=None), marked_image_file: Optional[UploadFile] = File(default=None), marked_image_url: Optional[str] = Form(default=None), marked_image_base64: Optional[str] = Form(default=None), original_image_file: Optional[UploadFile] = File(default=None), original_image_url: Optional[str] = Form(default=None), original_image_base64: Optional[str] = Form(default=None), ): if x_internal_token != password: return JSONResponse(status_code=401, content={"detail": "unauthorized"}) delay = float(request.query_params.get("delay", "1")) data = { "best_hairline_image_base64": f"data:image/png;base64,{TINY_PNG_BASE64}", "hair_growth_image_base64": f"data:image/png;base64,{TINY_PNG_BASE64}", "hairline_type": "花瓣形", } return mock_ok(data, delay=delay) # --- 接口4 --- @app.post("/api/v1/face/features") async def face_features( request: Request, x_internal_token: Optional[str] = Header(default=None), image_file: Optional[UploadFile] = File(default=None), image_url: Optional[str] = Form(default=None), image_base64: Optional[str] = Form(default=None), ): if x_internal_token != password: return JSONResponse(status_code=401, content={"detail": "unauthorized"}) delay = float(request.query_params.get("delay", "1")) features = json.dumps( { "face_shape": "鹅蛋脸", "eyebrow_shape": "柳叶眉", "facial_age": 26, "dynamic_static_type": "静态", "gender": "女", "gene_style": {"label": "面部特征标签", "style": "基因风格示例"}, }, ensure_ascii=False, ) return mock_ok({"features": features}, delay=delay) # --- 接口5 --- @app.post("/api/v1/hairline/generate") async def hairline_generate( request: Request, x_internal_token: Optional[str] = Header(default=None), image_file: Optional[UploadFile] = File(default=None), image_url: Optional[str] = Form(default=None), image_base64: Optional[str] = Form(default=None), ): if x_internal_token != password: return JSONResponse(status_code=401, content={"detail": "unauthorized"}) delay = float(request.query_params.get("delay", "1")) data = { "hairline_images": [ { "image_base64": f"data:image/png;base64,{TINY_PNG_BASE64}", "order": 1, }, { "image_base64": f"data:image/png;base64,{TINY_PNG_BASE64}", "order": 2, }, ], "best_hairline_center_point": {"x": 540, "y": 430}, } return mock_ok(data, delay=delay) return app # --------------------------------------------------------------------------- # 入口 # --------------------------------------------------------------------------- if __name__ == "__main__": parser = argparse.ArgumentParser(description="Stub Worker for gateway testing") parser.add_argument("--port", type=int, default=28187, help="监听端口（默认 28187）") parser.add_argument("--password", type=str, default="dev-shared-secret-2026", help="共享密码") args = parser.parse_args() app = create_app(port=args.port, password=args.password) print(f"[stub_worker] 启动在 :{args.port}，密码={args.password[:4]}...") uvicorn.run(app, host="127.0.0.1", port=args.port, log_level="warning")