xsl
14ec4964a0
请求头X-API-Token默认yonghe1999，可通过interface52_token覆盖；缺少或不匹配返回401/code1009，不进入派发队列。客户端Token不转发到worker，其他接口保持原样。 同步Token输入框、配置示例与调用文档。57项网关测试及临时网关连接真实worker的浏览器验证通过。
229 lines
10 KiB
Python
229 lines
10 KiB
Python
"""Gateway-only interface 52 integration: raw forwarding, token and image URLs."""
|
|
import asyncio
|
|
import base64
|
|
import importlib
|
|
import io
|
|
from types import SimpleNamespace
|
|
from unittest.mock import AsyncMock, Mock
|
|
|
|
import httpx
|
|
from PIL import Image
|
|
import pytest
|
|
from fastapi.testclient import TestClient
|
|
|
|
from gateway import config, forward
|
|
|
|
gateway = importlib.import_module("gateway.app")
|
|
PATH = "/api/v1/hairline/generate-52"
|
|
CLIENT_HEADERS = {"X-API-Token": "yonghe1999"}
|
|
|
|
|
|
@pytest.mark.parametrize("headers", [
|
|
{}, {"X-API-Token": ""}, {"X-API-Token": "wrong"},
|
|
{"X-API-Token": "yonghe1999 "}, {"X-Internal-Token": "yonghe1999"},
|
|
])
|
|
def test_invalid_client_token_is_rejected_before_dispatch(monkeypatch, headers):
|
|
monkeypatch.setattr(config, "get_config", lambda: {})
|
|
proxy = AsyncMock()
|
|
monkeypatch.setattr(gateway, "_proxy", proxy)
|
|
client = TestClient(gateway.app)
|
|
try:
|
|
response = client.post(PATH, headers=headers)
|
|
assert response.status_code == 401
|
|
body = response.json()
|
|
assert body["code"] == 1009 and body["data"] is None
|
|
assert body["request_id"].startswith("gw-")
|
|
finally:
|
|
client.close()
|
|
proxy.assert_not_awaited()
|
|
|
|
|
|
def test_custom_token_replaces_default_only_on_interface52(monkeypatch):
|
|
monkeypatch.setattr(config, "get_config", lambda: {"interface52_token": "custom-client-token"})
|
|
proxy = AsyncMock(return_value={"code": 0, "data": None})
|
|
monkeypatch.setattr(gateway, "_proxy", proxy)
|
|
client = TestClient(gateway.app)
|
|
try:
|
|
assert client.post(PATH, headers=CLIENT_HEADERS).status_code == 401
|
|
assert client.post(PATH, headers={"X-API-Token": "custom-client-token"}).status_code == 200
|
|
assert client.post("/api/v1/hairline/generate").status_code == 200
|
|
finally:
|
|
client.close()
|
|
assert proxy.await_count == 2
|
|
|
|
|
|
@pytest.mark.parametrize("token", ["", None, 123, []])
|
|
def test_invalid_configured_token_cannot_disable_auth(monkeypatch, token):
|
|
monkeypatch.setattr(config, "get_config", lambda: {"interface52_token": token})
|
|
proxy = AsyncMock()
|
|
monkeypatch.setattr(gateway, "_proxy", proxy)
|
|
client = TestClient(gateway.app)
|
|
try:
|
|
assert client.post(PATH, headers=CLIENT_HEADERS).status_code == 401
|
|
finally:
|
|
client.close()
|
|
proxy.assert_not_awaited()
|
|
|
|
|
|
@pytest.mark.parametrize("token", ["", None, 123, []])
|
|
def test_startup_rejects_invalid_client_token_configuration(tmp_path, monkeypatch, token):
|
|
import json
|
|
path = tmp_path / "config.json"
|
|
path.write_text(json.dumps({"interface52_token": token}))
|
|
monkeypatch.setattr(config, "CONFIG_PATH", path)
|
|
monkeypatch.setattr(config, "_config_cache", None)
|
|
with pytest.raises(ValueError, match="interface52_token"):
|
|
config.load_config()
|
|
|
|
|
|
def test_existing_config_gets_default_client_token(tmp_path, monkeypatch):
|
|
import json
|
|
path = tmp_path / "config.json"
|
|
path.write_text(json.dumps({"workers": ["http://worker"], "shared_password": "test-internal-token",
|
|
"static_dir": str(tmp_path)}))
|
|
monkeypatch.setattr(config, "CONFIG_PATH", path)
|
|
monkeypatch.setattr(config, "_config_cache", None)
|
|
assert config.load_config()["interface52_token"] == "yonghe1999"
|
|
|
|
|
|
def image_bytes(fmt):
|
|
buf = io.BytesIO()
|
|
Image.new("RGB", (6, 8), (127, 127, 127)).save(buf, format=fmt)
|
|
return buf.getvalue()
|
|
|
|
|
|
@pytest.fixture
|
|
def proxy_setup(tmp_path, monkeypatch):
|
|
output = tmp_path / "annotations"
|
|
output.mkdir()
|
|
cfg = {"shared_password": "private-gateway-token", "public_base_url": "http://testserver",
|
|
"static_dir": str(output), "dispatch": {"request_timeout_seconds": 600, "max_retries": 0}}
|
|
monkeypatch.setattr(config, "get_config", lambda: cfg)
|
|
acquire = AsyncMock(return_value=True)
|
|
release = Mock()
|
|
monkeypatch.setattr(forward, "acquire_global_slot", acquire)
|
|
monkeypatch.setattr(forward, "release_global_slot", release)
|
|
monkeypatch.setattr(forward, "acquire_worker", AsyncMock(return_value=SimpleNamespace(url="http://worker")))
|
|
release_worker = AsyncMock()
|
|
monkeypatch.setattr(forward, "release_worker", release_worker)
|
|
unhealthy = AsyncMock()
|
|
monkeypatch.setattr(forward, "mark_worker_unhealthy", unhealthy)
|
|
static = next(route.app for route in gateway.app.routes if getattr(route, "name", None) == "static")
|
|
monkeypatch.setattr(static, "all_directories", [str(tmp_path), *static.all_directories])
|
|
return cfg, acquire, release, release_worker, unhealthy
|
|
|
|
|
|
@pytest.mark.parametrize("path", ["/api/v1/hairline/generate", PATH])
|
|
def test_hairline_routes_forward_to_their_own_worker_path(monkeypatch, path):
|
|
monkeypatch.setattr(config, "get_config", lambda: {})
|
|
proxy = AsyncMock(return_value={"code": 0, "data": None})
|
|
monkeypatch.setattr(gateway, "_proxy", proxy)
|
|
client = TestClient(gateway.app, client=("203.0.113.9", 50000))
|
|
try:
|
|
assert client.post(path, headers=CLIENT_HEADERS if path == PATH else {}).json() == {"code": 0, "data": None}
|
|
finally:
|
|
client.close()
|
|
proxy.assert_awaited_once()
|
|
assert proxy.await_args.args[1] == path
|
|
|
|
|
|
@pytest.mark.parametrize("generate_grow_image", [False, True])
|
|
def test_interface52_raw_forwarding_and_served_image_urls(proxy_setup, monkeypatch, generate_grow_image):
|
|
cfg, acquire, release, release_worker, unhealthy = proxy_setup
|
|
png, jpg = image_bytes("PNG"), image_bytes("JPEG")
|
|
metadata = {"coordinate_space": "processed_image", "width": 6, "height": 8,
|
|
"framing": {"status": "reference_crop", "axis_x_error_px": 0},
|
|
"height_rule": {"target_center": {"x": 3.0, "y": 2.0}},
|
|
"guide_landmarks": {"9": {"x": 3.0, "y": 3.0}},
|
|
"best_hairline_center_point": {"x": 3, "y": 2},
|
|
"high_hairline_center_point": None, "low_hairline_center_point": None,
|
|
"face_measure": None}
|
|
worker_json = {"code": 0, "data": {**metadata,
|
|
"input_image_base64": base64.b64encode(jpg).decode(),
|
|
"original_image_base64": base64.b64encode(jpg).decode(),
|
|
"hairline_images": [{"hairline_type": "ellipse", "order": 1,
|
|
"height_placement": {"target_error_px": 0.01},
|
|
"image_middle_base64": base64.b64encode(png).decode(),
|
|
"image_high_base64": None, "image_low_base64": None,
|
|
"grown_image_base64": base64.b64encode(jpg).decode() if generate_grow_image else None}]}}
|
|
calls = []
|
|
|
|
def worker(request):
|
|
calls.append(request)
|
|
return httpx.Response(200, json=worker_json)
|
|
|
|
upstream = httpx.AsyncClient(transport=httpx.MockTransport(worker))
|
|
monkeypatch.setattr(forward, "get_client", lambda: upstream)
|
|
fields = {"gender": "male", "hair_style": "1,3", "normalize_image": "true",
|
|
"generate_grow_image": str(generate_grow_image).lower(), "use_mask": "false",
|
|
"prompt": "test prompt"}
|
|
source = httpx.Request("POST", "http://testserver" + PATH,
|
|
files={"image_file": ("face.png", png, "image/png")}, data=fields)
|
|
raw = source.read()
|
|
client = TestClient(gateway.app, client=("203.0.113.9", 50000))
|
|
try:
|
|
result = client.post(PATH, content=raw, headers={
|
|
**CLIENT_HEADERS,
|
|
"Content-Type": source.headers["content-type"], "X-Internal-Token": "untrusted-client-token"})
|
|
assert result.status_code == 200
|
|
data = result.json()["data"]
|
|
assert {key: data[key] for key in metadata} == metadata
|
|
item = data["hairline_images"][0]
|
|
assert item["image_high_base64"] is None and item["image_low_base64"] is None
|
|
assert "image_high_url" not in item and "image_low_url" not in item
|
|
assert item["height_placement"] == {"target_error_px": 0.01}
|
|
for key, expected in [("input_image_url", jpg), ("original_image_url", jpg)]:
|
|
assert client.get(data[key]).content == expected
|
|
assert data[key].endswith(".jpg")
|
|
assert client.get(item["image_middle_url"]).content == png
|
|
assert item["image_middle_url"].endswith(".png")
|
|
if generate_grow_image:
|
|
assert client.get(item["grown_image_url"]).content == jpg
|
|
else:
|
|
assert item["grown_image_base64"] is None and "grown_image_url" not in item
|
|
finally:
|
|
client.close()
|
|
asyncio.run(upstream.aclose())
|
|
assert len(calls) == 1
|
|
assert str(calls[0].url) == "http://worker" + PATH
|
|
assert calls[0].content == raw
|
|
assert calls[0].headers["content-type"] == source.headers["content-type"]
|
|
assert calls[0].headers["x-internal-token"] == cfg["shared_password"]
|
|
assert calls[0].headers.get("x-forwarded-for") is None
|
|
assert calls[0].headers.get("x-api-token") is None
|
|
acquire.assert_awaited_once_with(cfg, 0)
|
|
release.assert_called_once()
|
|
release_worker.assert_awaited_once()
|
|
unhealthy.assert_not_awaited()
|
|
|
|
|
|
def test_worker_ip_rejection_is_preserved(proxy_setup, monkeypatch):
|
|
_, _, release, release_worker, unhealthy = proxy_setup
|
|
body = {"code": 1010, "message": "来源IP不在接口白名单", "data": None}
|
|
upstream = httpx.AsyncClient(transport=httpx.MockTransport(lambda request: httpx.Response(403, json=body)))
|
|
monkeypatch.setattr(forward, "get_client", lambda: upstream)
|
|
client = TestClient(gateway.app)
|
|
try:
|
|
response = client.post(PATH, headers=CLIENT_HEADERS, data={"gender": "male", "hair_style": "1"})
|
|
assert response.status_code == 403 and response.json() == body
|
|
finally:
|
|
client.close()
|
|
asyncio.run(upstream.aclose())
|
|
release.assert_called_once()
|
|
release_worker.assert_awaited_once()
|
|
unhealthy.assert_not_awaited()
|
|
|
|
|
|
def test_public_interface52_page_has_no_internal_password():
|
|
client = TestClient(gateway.app)
|
|
try:
|
|
page = client.get("/static/test_interface52.html")
|
|
assert page.status_code == 200
|
|
assert PATH in page.text and "X-Internal-Token" not in page.text
|
|
assert 'id="apiToken"' in page.text and "X-API-Token" in page.text
|
|
assert "localhost:8187" not in page.text
|
|
assert "test_interface52_compare.html" not in page.text
|
|
assert client.get("/static/hairline_reference_box.js").status_code == 200
|
|
finally:
|
|
client.close()
|