257 lines
9.7 KiB
Python
257 lines
9.7 KiB
Python
"""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")
|