Spaces:
Running
Running
Keep HEAD metadata consistent and update existing API contract tests
Browse files- job_api_common.py +4 -0
- tests/test_image_inversion_job_api.py +775 -773
job_api_common.py
CHANGED
|
@@ -21,6 +21,10 @@ class ResultFileResponse(FileResponse):
|
|
| 21 |
"""
|
| 22 |
release = self.background
|
| 23 |
self.background = None
|
|
|
|
|
|
|
|
|
|
|
|
|
| 24 |
try:
|
| 25 |
await super().__call__(scope, receive, send)
|
| 26 |
finally:
|
|
|
|
| 21 |
"""
|
| 22 |
release = self.background
|
| 23 |
self.background = None
|
| 24 |
+
if scope["method"] == "HEAD":
|
| 25 |
+
# Range 只用于 GET;HEAD 应提供完整结果的长度和校验信息。
|
| 26 |
+
scope = {**scope, "headers": [(name, value) for name, value in scope["headers"]
|
| 27 |
+
if name.lower() != b"range"]}
|
| 28 |
try:
|
| 29 |
await super().__call__(scope, receive, send)
|
| 30 |
finally:
|
tests/test_image_inversion_job_api.py
CHANGED
|
@@ -1,773 +1,775 @@
|
|
| 1 |
-
from __future__ import annotations
|
| 2 |
-
|
| 3 |
-
import ast
|
| 4 |
-
import io
|
| 5 |
-
import logging
|
| 6 |
-
import os
|
| 7 |
-
import socket
|
| 8 |
-
import tempfile
|
| 9 |
-
import threading
|
| 10 |
-
import time
|
| 11 |
-
import unittest
|
| 12 |
-
from concurrent.futures import ThreadPoolExecutor
|
| 13 |
-
from pathlib import Path
|
| 14 |
-
from unittest.mock import patch
|
| 15 |
-
|
| 16 |
-
import httpx
|
| 17 |
-
from fastapi import FastAPI
|
| 18 |
-
from fastapi.testclient import TestClient
|
| 19 |
-
from PIL import Image
|
| 20 |
-
from pydantic import ValidationError
|
| 21 |
-
|
| 22 |
-
from image_inversion_job_api import (
|
| 23 |
-
ImageInversionJobAPI,
|
| 24 |
-
ImageInversionJobRecord,
|
| 25 |
-
ImageInversionJobRequest,
|
| 26 |
-
ImageInversionJobSettings,
|
| 27 |
-
ImageSourceError,
|
| 28 |
-
create_job_api_lifespan,
|
| 29 |
-
decode_image_bytes,
|
| 30 |
-
fetch_remote_image,
|
| 31 |
-
fingerprint_request,
|
| 32 |
-
validate_image_url,
|
| 33 |
-
)
|
| 34 |
-
|
| 35 |
-
|
| 36 |
-
API_KEY = "k" * 48
|
| 37 |
-
EXPECTED_RESULT = {
|
| 38 |
-
"status_markdown": "done",
|
| 39 |
-
"general_tags_html": "<p>general</p>",
|
| 40 |
-
"character_tags_html": "<p>character</p>",
|
| 41 |
-
"ip_tags_html": "<p>ip</p>",
|
| 42 |
-
"summary_text": "tag one, tag two",
|
| 43 |
-
"metadata": {"device": "cpu", "latency_s_total": 0.1},
|
| 44 |
-
}
|
| 45 |
-
|
| 46 |
-
|
| 47 |
-
def _make_png_bytes(size: tuple[int, int] = (96, 80)) -> bytes:
|
| 48 |
-
"""生成远程抓图测试使用的有效 PNG。
|
| 49 |
-
|
| 50 |
-
Args:
|
| 51 |
-
size: PNG 的宽高。
|
| 52 |
-
|
| 53 |
-
Returns:
|
| 54 |
-
完整 PNG 文件字节。
|
| 55 |
-
"""
|
| 56 |
-
image_buffer = io.BytesIO()
|
| 57 |
-
with Image.new("RGB", size, "green") as image:
|
| 58 |
-
image.save(image_buffer, format="PNG")
|
| 59 |
-
return image_buffer.getvalue()
|
| 60 |
-
|
| 61 |
-
|
| 62 |
-
PNG_BYTES = _make_png_bytes()
|
| 63 |
-
|
| 64 |
-
|
| 65 |
-
class _TrackedImage:
|
| 66 |
-
"""记录任务输入图是否被 API 生命周期可靠关闭。"""
|
| 67 |
-
|
| 68 |
-
def __init__(self) -> None:
|
| 69 |
-
self.closed = False
|
| 70 |
-
|
| 71 |
-
def close(self) -> None:
|
| 72 |
-
"""标记当前测试图片已经关闭。"""
|
| 73 |
-
self.closed = True
|
| 74 |
-
|
| 75 |
-
|
| 76 |
-
class _FakeStreamResponse:
|
| 77 |
-
"""提供 httpx 流式响应所需的最小测试替身。"""
|
| 78 |
-
|
| 79 |
-
def __init__(
|
| 80 |
-
self,
|
| 81 |
-
status_code: int,
|
| 82 |
-
body: bytes,
|
| 83 |
-
*,
|
| 84 |
-
peer_address: str = "93.184.216.34",
|
| 85 |
-
extra_headers: dict[str, str] | None = None,
|
| 86 |
-
) -> None:
|
| 87 |
-
self.status_code = status_code
|
| 88 |
-
self.body = body
|
| 89 |
-
self.headers = {
|
| 90 |
-
"content-type": "image/png",
|
| 91 |
-
"content-length": str(len(body)),
|
| 92 |
-
}
|
| 93 |
-
if extra_headers:
|
| 94 |
-
self.headers.update(extra_headers)
|
| 95 |
-
self.extensions = {
|
| 96 |
-
"network_stream": type(
|
| 97 |
-
"NetworkStreamStub",
|
| 98 |
-
(),
|
| 99 |
-
{
|
| 100 |
-
"get_extra_info": lambda self, name: (
|
| 101 |
-
(peer_address, 443) if name == "server_addr" else None
|
| 102 |
-
)
|
| 103 |
-
},
|
| 104 |
-
)()
|
| 105 |
-
}
|
| 106 |
-
|
| 107 |
-
def __enter__(self) -> "_FakeStreamResponse":
|
| 108 |
-
"""进入流式响应上下文并返回自身。"""
|
| 109 |
-
return self
|
| 110 |
-
|
| 111 |
-
def __exit__(self, exc_type, exc, traceback) -> bool:
|
| 112 |
-
"""退出响应上下文且不吞掉异常。"""
|
| 113 |
-
del exc_type, exc, traceback
|
| 114 |
-
return False
|
| 115 |
-
|
| 116 |
-
def iter_bytes(self):
|
| 117 |
-
"""以单个分块返回完整响应体。"""
|
| 118 |
-
yield self.body
|
| 119 |
-
|
| 120 |
-
|
| 121 |
-
class ImageInversionJobAPITest(unittest.TestCase):
|
| 122 |
-
"""验证 Image Inversion 自定义任务 API 的公开契约与安全边界。"""
|
| 123 |
-
|
| 124 |
-
def setUp(self) -> None:
|
| 125 |
-
"""登记每个测试需要在结束时关闭的客户端和线程事件。"""
|
| 126 |
-
self.clients: list[TestClient] = []
|
| 127 |
-
self.release_events: list[threading.Event] = []
|
| 128 |
-
|
| 129 |
-
def tearDown(self) -> None:
|
| 130 |
-
"""释放测试线程和 TestClient,不创建持久化测试产物。"""
|
| 131 |
-
for release_event in self.release_events:
|
| 132 |
-
release_event.set()
|
| 133 |
-
time.sleep(0.02)
|
| 134 |
-
for client in self.clients:
|
| 135 |
-
client.close()
|
| 136 |
-
|
| 137 |
-
@staticmethod
|
| 138 |
-
def _settings(**overrides) -> ImageInversionJobSettings:
|
| 139 |
-
"""构造不依赖进程环境的隔离任务设置。
|
| 140 |
-
|
| 141 |
-
Args:
|
| 142 |
-
overrides: 要覆盖的 ImageInversionJobSettings 字段。
|
| 143 |
-
|
| 144 |
-
Returns:
|
| 145 |
-
可供单元测试直接注入的设置。
|
| 146 |
-
"""
|
| 147 |
-
values = {
|
| 148 |
-
"api_key": API_KEY,
|
| 149 |
-
"allowed_hosts": frozenset({"allowed.example"}),
|
| 150 |
-
"result_ttl_seconds": 1800,
|
| 151 |
-
"poll_after_seconds": 1,
|
| 152 |
-
"fetch_timeout_seconds": 15.0,
|
| 153 |
-
"max_image_bytes": 8 * 1024 * 1024,
|
| 154 |
-
"max_image_pixels": 4_194_304,
|
| 155 |
-
"max_image_dimension": 4096,
|
| 156 |
-
"max_records": 256,
|
| 157 |
-
"space_host": "https://space.example",
|
| 158 |
-
}
|
| 159 |
-
values.update(overrides)
|
| 160 |
-
return ImageInversionJobSettings(**values)
|
| 161 |
-
|
| 162 |
-
@staticmethod
|
| 163 |
-
def _payload(**overrides) -> dict:
|
| 164 |
-
"""生成可直接提交给创建接口的合法请求体。
|
| 165 |
-
|
| 166 |
-
Args:
|
| 167 |
-
overrides: 要覆盖或增加的 JSON 字段。
|
| 168 |
-
|
| 169 |
-
Returns:
|
| 170 |
-
包含远程输入图和全部展示参数的请求字典。
|
| 171 |
-
"""
|
| 172 |
-
payload = {
|
| 173 |
-
"input_image_url": "https://allowed.example/input.png?source=private",
|
| 174 |
-
"general_threshold": 0.35,
|
| 175 |
-
"character_threshold": 0.85,
|
| 176 |
-
"show_confidence": True,
|
| 177 |
-
"show_general": True,
|
| 178 |
-
"show_character": True,
|
| 179 |
-
"show_ip": True,
|
| 180 |
-
"separator": "comma",
|
| 181 |
-
"show_chinese": True,
|
| 182 |
-
}
|
| 183 |
-
payload.update(overrides)
|
| 184 |
-
return payload
|
| 185 |
-
|
| 186 |
-
@staticmethod
|
| 187 |
-
def _headers(**overrides) -> dict[str, str]:
|
| 188 |
-
"""生成包含共享密钥的请求头。
|
| 189 |
-
|
| 190 |
-
Args:
|
| 191 |
-
overrides: 要覆盖或追加的请求头。
|
| 192 |
-
|
| 193 |
-
Returns:
|
| 194 |
-
可传给 TestClient 的请求头字典。
|
| 195 |
-
"""
|
| 196 |
-
headers = {"X-API-Key": API_KEY}
|
| 197 |
-
headers.update(overrides)
|
| 198 |
-
return headers
|
| 199 |
-
|
| 200 |
-
def _make_service(
|
| 201 |
-
self,
|
| 202 |
-
*,
|
| 203 |
-
executor=None,
|
| 204 |
-
image_fetcher=None,
|
| 205 |
-
inference_slot: threading.Lock | None = None,
|
| 206 |
-
settings: ImageInversionJobSettings | None = None,
|
| 207 |
-
add_catch_all: bool = False,
|
| 208 |
-
) -> tuple[TestClient, ImageInversionJobAPI, list[_TrackedImage], dict[str, int]]:
|
| 209 |
-
"""创建完全隔离且不导入模型的 FastAPI 测试服务。
|
| 210 |
-
|
| 211 |
-
Args:
|
| 212 |
-
executor: 可选假标签执行器。
|
| 213 |
-
image_fetcher: 可选假远程图片抓取器。
|
| 214 |
-
inference_slot: 可选 UI/API 共享锁。
|
| 215 |
-
settings: 可选任务设置。
|
| 216 |
-
add_catch_all: 是否先注册模拟 Gradio 的宽泛路由。
|
| 217 |
-
|
| 218 |
-
Returns:
|
| 219 |
-
TestClient、API 对象、输入图片列表与调用计数器。
|
| 220 |
-
"""
|
| 221 |
-
settings = settings or self._settings()
|
| 222 |
-
created_images: list[_TrackedImage] = []
|
| 223 |
-
counters = {"fetch": 0, "execute": 0}
|
| 224 |
-
counter_lock = threading.Lock()
|
| 225 |
-
|
| 226 |
-
if image_fetcher is None:
|
| 227 |
-
|
| 228 |
-
def image_fetcher(url, current_settings):
|
| 229 |
-
"""返回可追踪关闭状态的假输入图。"""
|
| 230 |
-
del url, current_settings
|
| 231 |
-
with counter_lock:
|
| 232 |
-
counters["fetch"] += 1
|
| 233 |
-
image = _TrackedImage()
|
| 234 |
-
created_images.append(image)
|
| 235 |
-
return image
|
| 236 |
-
|
| 237 |
-
if executor is None:
|
| 238 |
-
|
| 239 |
-
def executor(payload, input_image, job_id):
|
| 240 |
-
"""返回与新接口契约一致的六字段 JSON 结果。"""
|
| 241 |
-
del payload, input_image, job_id
|
| 242 |
-
with counter_lock:
|
| 243 |
-
counters["execute"] += 1
|
| 244 |
-
return EXPECTED_RESULT
|
| 245 |
-
|
| 246 |
-
api = ImageInversionJobAPI(
|
| 247 |
-
settings=settings,
|
| 248 |
-
executor=executor,
|
| 249 |
-
inference_slot=inference_slot or threading.Lock(),
|
| 250 |
-
image_fetcher=image_fetcher,
|
| 251 |
-
)
|
| 252 |
-
app = FastAPI()
|
| 253 |
-
if add_catch_all:
|
| 254 |
-
|
| 255 |
-
@app.api_route("/{path:path}", methods=["GET", "POST"])
|
| 256 |
-
def catch_all(path: str):
|
| 257 |
-
"""模拟可能抢先匹配 API 请求的 Gradio 宽泛路由。"""
|
| 258 |
-
return {"catch_all": path}
|
| 259 |
-
|
| 260 |
-
api.install_on_app(app)
|
| 261 |
-
client = TestClient(app, base_url="https://space.example")
|
| 262 |
-
self.clients.append(client)
|
| 263 |
-
return client, api, created_images, counters
|
| 264 |
-
|
| 265 |
-
def _wait_for_terminal(
|
| 266 |
-
self,
|
| 267 |
-
client: TestClient,
|
| 268 |
-
job_id: str,
|
| 269 |
-
timeout_seconds: float = 3.0,
|
| 270 |
-
):
|
| 271 |
-
"""轮询测试任务直到返回成功或失败终态。
|
| 272 |
-
|
| 273 |
-
Args:
|
| 274 |
-
client: 当前隔离服务客户端。
|
| 275 |
-
job_id: 创建接口返回的任务标识。
|
| 276 |
-
timeout_seconds: 允许后台测试线程运行的最长时间。
|
| 277 |
-
|
| 278 |
-
Returns:
|
| 279 |
-
第一个非 202 状态响应。
|
| 280 |
-
"""
|
| 281 |
-
deadline = time.monotonic() + timeout_seconds
|
| 282 |
-
while time.monotonic() < deadline:
|
| 283 |
-
response = client.get(
|
| 284 |
-
f"/api/jobs/{job_id}",
|
| 285 |
-
headers={"X-API-Key": API_KEY},
|
| 286 |
-
)
|
| 287 |
-
if response.status_code != 202:
|
| 288 |
-
return response
|
| 289 |
-
time.sleep(0.01)
|
| 290 |
-
self.fail(f"job {job_id} did not reach a terminal state")
|
| 291 |
-
|
| 292 |
-
def test_auth_schema_cpu_contract_and_route_precedence(self) -> None:
|
| 293 |
-
"""验证鉴权、CPU 请求契约及自定义路由优先级。"""
|
| 294 |
-
client, api, _, counters = self._make_service(add_catch_all=True)
|
| 295 |
-
self.assertTrue(callable(create_job_api_lifespan(api)))
|
| 296 |
-
|
| 297 |
-
unauthorized = client.post("/api/jobs", json=self._payload())
|
| 298 |
-
self.assertEqual(unauthorized.status_code, 401)
|
| 299 |
-
self.assertEqual(counters["fetch"], 0)
|
| 300 |
-
|
| 301 |
-
unknown_field = client.post(
|
| 302 |
-
"/api/jobs",
|
| 303 |
-
json=self._payload(unknown=True),
|
| 304 |
-
headers=self._headers(),
|
| 305 |
-
)
|
| 306 |
-
self.assertEqual(unknown_field.status_code, 422)
|
| 307 |
-
self.assertEqual(counters["fetch"], 0)
|
| 308 |
-
|
| 309 |
-
# cpu-basic Space 不依赖 ZeroGPU 的 x-ip-token。
|
| 310 |
-
accepted = client.post(
|
| 311 |
-
"/api/jobs",
|
| 312 |
-
json=self._payload(),
|
| 313 |
-
headers=self._headers(),
|
| 314 |
-
)
|
| 315 |
-
self.assertEqual(accepted.status_code, 202)
|
| 316 |
-
self.assertNotIn("catch_all", accepted.json())
|
| 317 |
-
terminal = self._wait_for_terminal(client, accepted.json()["job_id"])
|
| 318 |
-
self.assertEqual(terminal.status_code, 200)
|
| 319 |
-
self.assertEqual(terminal.json(), EXPECTED_RESULT)
|
| 320 |
-
|
| 321 |
-
|
| 322 |
-
|
| 323 |
-
|
| 324 |
-
|
| 325 |
-
|
| 326 |
-
|
| 327 |
-
|
| 328 |
-
|
| 329 |
-
|
| 330 |
-
|
| 331 |
-
|
| 332 |
-
|
| 333 |
-
|
| 334 |
-
|
| 335 |
-
|
| 336 |
-
|
| 337 |
-
|
| 338 |
-
|
| 339 |
-
|
| 340 |
-
|
| 341 |
-
|
| 342 |
-
self.assertEqual(created.
|
| 343 |
-
|
| 344 |
-
|
| 345 |
-
|
| 346 |
-
self.
|
| 347 |
-
self.assertEqual(
|
| 348 |
-
self.assertEqual(
|
| 349 |
-
self.
|
| 350 |
-
self.
|
| 351 |
-
|
| 352 |
-
|
| 353 |
-
|
| 354 |
-
|
| 355 |
-
|
| 356 |
-
|
| 357 |
-
|
| 358 |
-
|
| 359 |
-
|
| 360 |
-
|
| 361 |
-
|
| 362 |
-
|
| 363 |
-
|
| 364 |
-
|
| 365 |
-
|
| 366 |
-
|
| 367 |
-
|
| 368 |
-
|
| 369 |
-
|
| 370 |
-
|
| 371 |
-
|
| 372 |
-
|
| 373 |
-
self.
|
| 374 |
-
|
| 375 |
-
|
| 376 |
-
|
| 377 |
-
|
| 378 |
-
|
| 379 |
-
self.assertEqual(
|
| 380 |
-
|
| 381 |
-
|
| 382 |
-
|
| 383 |
-
|
| 384 |
-
|
| 385 |
-
|
| 386 |
-
|
| 387 |
-
|
| 388 |
-
|
| 389 |
-
|
| 390 |
-
|
| 391 |
-
|
| 392 |
-
|
| 393 |
-
|
| 394 |
-
|
| 395 |
-
self.assertEqual(
|
| 396 |
-
|
| 397 |
-
|
| 398 |
-
|
| 399 |
-
|
| 400 |
-
self.
|
| 401 |
-
|
| 402 |
-
|
| 403 |
-
|
| 404 |
-
|
| 405 |
-
|
| 406 |
-
|
| 407 |
-
|
| 408 |
-
|
| 409 |
-
|
| 410 |
-
|
| 411 |
-
|
| 412 |
-
|
| 413 |
-
|
| 414 |
-
|
| 415 |
-
|
| 416 |
-
|
| 417 |
-
|
| 418 |
-
|
| 419 |
-
|
| 420 |
-
|
| 421 |
-
|
| 422 |
-
|
| 423 |
-
|
| 424 |
-
|
| 425 |
-
|
| 426 |
-
|
| 427 |
-
|
| 428 |
-
|
| 429 |
-
|
| 430 |
-
|
| 431 |
-
|
| 432 |
-
|
| 433 |
-
|
| 434 |
-
|
| 435 |
-
|
| 436 |
-
|
| 437 |
-
|
| 438 |
-
|
| 439 |
-
|
| 440 |
-
|
| 441 |
-
self.assertEqual(
|
| 442 |
-
self.
|
| 443 |
-
self.assertEqual(
|
| 444 |
-
|
| 445 |
-
|
| 446 |
-
|
| 447 |
-
|
| 448 |
-
|
| 449 |
-
|
| 450 |
-
|
| 451 |
-
|
| 452 |
-
|
| 453 |
-
|
| 454 |
-
|
| 455 |
-
|
| 456 |
-
|
| 457 |
-
|
| 458 |
-
|
| 459 |
-
|
| 460 |
-
|
| 461 |
-
|
| 462 |
-
|
| 463 |
-
|
| 464 |
-
|
| 465 |
-
|
| 466 |
-
|
| 467 |
-
|
| 468 |
-
|
| 469 |
-
self.
|
| 470 |
-
|
| 471 |
-
self.
|
| 472 |
-
self.
|
| 473 |
-
|
| 474 |
-
self.
|
| 475 |
-
|
| 476 |
-
|
| 477 |
-
|
| 478 |
-
|
| 479 |
-
|
| 480 |
-
|
| 481 |
-
|
| 482 |
-
|
| 483 |
-
|
| 484 |
-
|
| 485 |
-
|
| 486 |
-
|
| 487 |
-
|
| 488 |
-
|
| 489 |
-
|
| 490 |
-
|
| 491 |
-
|
| 492 |
-
self.
|
| 493 |
-
self.
|
| 494 |
-
|
| 495 |
-
|
| 496 |
-
|
| 497 |
-
|
| 498 |
-
|
| 499 |
-
|
| 500 |
-
|
| 501 |
-
|
| 502 |
-
|
| 503 |
-
|
| 504 |
-
|
| 505 |
-
|
| 506 |
-
|
| 507 |
-
|
| 508 |
-
|
| 509 |
-
|
| 510 |
-
|
| 511 |
-
|
| 512 |
-
|
| 513 |
-
|
| 514 |
-
|
| 515 |
-
|
| 516 |
-
|
| 517 |
-
|
| 518 |
-
|
| 519 |
-
|
| 520 |
-
|
| 521 |
-
|
| 522 |
-
|
| 523 |
-
|
| 524 |
-
|
| 525 |
-
|
| 526 |
-
|
| 527 |
-
|
| 528 |
-
|
| 529 |
-
|
| 530 |
-
|
| 531 |
-
|
| 532 |
-
|
| 533 |
-
|
| 534 |
-
|
| 535 |
-
|
| 536 |
-
|
| 537 |
-
|
| 538 |
-
|
| 539 |
-
|
| 540 |
-
|
| 541 |
-
|
| 542 |
-
|
| 543 |
-
|
| 544 |
-
|
| 545 |
-
|
| 546 |
-
|
| 547 |
-
|
| 548 |
-
|
| 549 |
-
|
| 550 |
-
self.
|
| 551 |
-
self.
|
| 552 |
-
self.
|
| 553 |
-
|
| 554 |
-
|
| 555 |
-
|
| 556 |
-
|
| 557 |
-
|
| 558 |
-
|
| 559 |
-
|
| 560 |
-
|
| 561 |
-
|
| 562 |
-
|
| 563 |
-
|
| 564 |
-
|
| 565 |
-
|
| 566 |
-
|
| 567 |
-
|
| 568 |
-
|
| 569 |
-
|
| 570 |
-
|
| 571 |
-
self.
|
| 572 |
-
self.assertEqual(
|
| 573 |
-
|
| 574 |
-
|
| 575 |
-
|
| 576 |
-
|
| 577 |
-
|
| 578 |
-
|
| 579 |
-
|
| 580 |
-
|
| 581 |
-
|
| 582 |
-
|
| 583 |
-
|
| 584 |
-
|
| 585 |
-
|
| 586 |
-
settings.
|
| 587 |
-
|
| 588 |
-
|
| 589 |
-
|
| 590 |
-
|
| 591 |
-
|
| 592 |
-
|
| 593 |
-
|
| 594 |
-
|
| 595 |
-
|
| 596 |
-
|
| 597 |
-
|
| 598 |
-
"
|
| 599 |
-
|
| 600 |
-
|
| 601 |
-
|
| 602 |
-
|
| 603 |
-
|
| 604 |
-
|
| 605 |
-
|
| 606 |
-
|
| 607 |
-
|
| 608 |
-
|
| 609 |
-
|
| 610 |
-
|
| 611 |
-
|
| 612 |
-
|
| 613 |
-
|
| 614 |
-
|
| 615 |
-
|
| 616 |
-
|
| 617 |
-
|
| 618 |
-
|
| 619 |
-
|
| 620 |
-
|
| 621 |
-
|
| 622 |
-
|
| 623 |
-
|
| 624 |
-
|
| 625 |
-
|
| 626 |
-
self._payload(
|
| 627 |
-
self._payload(general_threshold=
|
| 628 |
-
self._payload(
|
| 629 |
-
self._payload(
|
| 630 |
-
|
| 631 |
-
|
| 632 |
-
|
| 633 |
-
|
| 634 |
-
|
| 635 |
-
|
| 636 |
-
|
| 637 |
-
|
| 638 |
-
|
| 639 |
-
|
| 640 |
-
|
| 641 |
-
socket.
|
| 642 |
-
|
| 643 |
-
|
| 644 |
-
|
| 645 |
-
|
| 646 |
-
|
| 647 |
-
|
| 648 |
-
|
| 649 |
-
|
| 650 |
-
|
| 651 |
-
|
| 652 |
-
|
| 653 |
-
|
| 654 |
-
|
| 655 |
-
|
| 656 |
-
|
| 657 |
-
"
|
| 658 |
-
"https://
|
| 659 |
-
"https://allowed.example/image.png
|
| 660 |
-
|
| 661 |
-
|
| 662 |
-
|
| 663 |
-
|
| 664 |
-
|
| 665 |
-
|
| 666 |
-
|
| 667 |
-
|
| 668 |
-
socket.
|
| 669 |
-
|
| 670 |
-
|
| 671 |
-
|
| 672 |
-
|
| 673 |
-
|
| 674 |
-
|
| 675 |
-
|
| 676 |
-
|
| 677 |
-
|
| 678 |
-
|
| 679 |
-
|
| 680 |
-
|
| 681 |
-
decoded
|
| 682 |
-
|
| 683 |
-
|
| 684 |
-
|
| 685 |
-
|
| 686 |
-
|
| 687 |
-
|
| 688 |
-
|
| 689 |
-
|
| 690 |
-
|
| 691 |
-
|
| 692 |
-
|
| 693 |
-
|
| 694 |
-
|
| 695 |
-
|
| 696 |
-
|
| 697 |
-
|
| 698 |
-
|
| 699 |
-
|
| 700 |
-
|
| 701 |
-
|
| 702 |
-
|
| 703 |
-
|
| 704 |
-
|
| 705 |
-
|
| 706 |
-
|
| 707 |
-
|
| 708 |
-
captured["
|
| 709 |
-
|
| 710 |
-
|
| 711 |
-
|
| 712 |
-
|
| 713 |
-
|
| 714 |
-
|
| 715 |
-
|
| 716 |
-
|
| 717 |
-
|
| 718 |
-
|
| 719 |
-
|
| 720 |
-
|
| 721 |
-
|
| 722 |
-
|
| 723 |
-
self.
|
| 724 |
-
|
| 725 |
-
|
| 726 |
-
|
| 727 |
-
|
| 728 |
-
|
| 729 |
-
|
| 730 |
-
|
| 731 |
-
|
| 732 |
-
|
| 733 |
-
|
| 734 |
-
|
| 735 |
-
|
| 736 |
-
|
| 737 |
-
|
| 738 |
-
|
| 739 |
-
|
| 740 |
-
|
| 741 |
-
|
| 742 |
-
|
| 743 |
-
|
| 744 |
-
|
| 745 |
-
|
| 746 |
-
|
| 747 |
-
|
| 748 |
-
|
| 749 |
-
|
| 750 |
-
|
| 751 |
-
|
| 752 |
-
|
| 753 |
-
|
| 754 |
-
|
| 755 |
-
|
| 756 |
-
|
| 757 |
-
|
| 758 |
-
|
| 759 |
-
|
| 760 |
-
|
| 761 |
-
|
| 762 |
-
|
| 763 |
-
|
| 764 |
-
|
| 765 |
-
|
| 766 |
-
self.
|
| 767 |
-
self.assertIn(
|
| 768 |
-
self.assertIn('
|
| 769 |
-
self.assertIn('
|
| 770 |
-
|
| 771 |
-
|
| 772 |
-
|
| 773 |
-
|
|
|
|
|
|
|
|
|
| 1 |
+
from __future__ import annotations
|
| 2 |
+
|
| 3 |
+
import ast
|
| 4 |
+
import io
|
| 5 |
+
import logging
|
| 6 |
+
import os
|
| 7 |
+
import socket
|
| 8 |
+
import tempfile
|
| 9 |
+
import threading
|
| 10 |
+
import time
|
| 11 |
+
import unittest
|
| 12 |
+
from concurrent.futures import ThreadPoolExecutor
|
| 13 |
+
from pathlib import Path
|
| 14 |
+
from unittest.mock import patch
|
| 15 |
+
|
| 16 |
+
import httpx
|
| 17 |
+
from fastapi import FastAPI
|
| 18 |
+
from fastapi.testclient import TestClient
|
| 19 |
+
from PIL import Image
|
| 20 |
+
from pydantic import ValidationError
|
| 21 |
+
|
| 22 |
+
from image_inversion_job_api import (
|
| 23 |
+
ImageInversionJobAPI,
|
| 24 |
+
ImageInversionJobRecord,
|
| 25 |
+
ImageInversionJobRequest,
|
| 26 |
+
ImageInversionJobSettings,
|
| 27 |
+
ImageSourceError,
|
| 28 |
+
create_job_api_lifespan,
|
| 29 |
+
decode_image_bytes,
|
| 30 |
+
fetch_remote_image,
|
| 31 |
+
fingerprint_request,
|
| 32 |
+
validate_image_url,
|
| 33 |
+
)
|
| 34 |
+
|
| 35 |
+
|
| 36 |
+
API_KEY = "k" * 48
|
| 37 |
+
EXPECTED_RESULT = {
|
| 38 |
+
"status_markdown": "done",
|
| 39 |
+
"general_tags_html": "<p>general</p>",
|
| 40 |
+
"character_tags_html": "<p>character</p>",
|
| 41 |
+
"ip_tags_html": "<p>ip</p>",
|
| 42 |
+
"summary_text": "tag one, tag two",
|
| 43 |
+
"metadata": {"device": "cpu", "latency_s_total": 0.1},
|
| 44 |
+
}
|
| 45 |
+
|
| 46 |
+
|
| 47 |
+
def _make_png_bytes(size: tuple[int, int] = (96, 80)) -> bytes:
|
| 48 |
+
"""生成远程抓图测试使用的有效 PNG。
|
| 49 |
+
|
| 50 |
+
Args:
|
| 51 |
+
size: PNG 的宽高。
|
| 52 |
+
|
| 53 |
+
Returns:
|
| 54 |
+
完整 PNG 文件字节。
|
| 55 |
+
"""
|
| 56 |
+
image_buffer = io.BytesIO()
|
| 57 |
+
with Image.new("RGB", size, "green") as image:
|
| 58 |
+
image.save(image_buffer, format="PNG")
|
| 59 |
+
return image_buffer.getvalue()
|
| 60 |
+
|
| 61 |
+
|
| 62 |
+
PNG_BYTES = _make_png_bytes()
|
| 63 |
+
|
| 64 |
+
|
| 65 |
+
class _TrackedImage:
|
| 66 |
+
"""记录任务输入图是否被 API 生命周期可靠关闭。"""
|
| 67 |
+
|
| 68 |
+
def __init__(self) -> None:
|
| 69 |
+
self.closed = False
|
| 70 |
+
|
| 71 |
+
def close(self) -> None:
|
| 72 |
+
"""标记当前测试图片已经关闭。"""
|
| 73 |
+
self.closed = True
|
| 74 |
+
|
| 75 |
+
|
| 76 |
+
class _FakeStreamResponse:
|
| 77 |
+
"""提供 httpx 流式响应所需的最小测试替身。"""
|
| 78 |
+
|
| 79 |
+
def __init__(
|
| 80 |
+
self,
|
| 81 |
+
status_code: int,
|
| 82 |
+
body: bytes,
|
| 83 |
+
*,
|
| 84 |
+
peer_address: str = "93.184.216.34",
|
| 85 |
+
extra_headers: dict[str, str] | None = None,
|
| 86 |
+
) -> None:
|
| 87 |
+
self.status_code = status_code
|
| 88 |
+
self.body = body
|
| 89 |
+
self.headers = {
|
| 90 |
+
"content-type": "image/png",
|
| 91 |
+
"content-length": str(len(body)),
|
| 92 |
+
}
|
| 93 |
+
if extra_headers:
|
| 94 |
+
self.headers.update(extra_headers)
|
| 95 |
+
self.extensions = {
|
| 96 |
+
"network_stream": type(
|
| 97 |
+
"NetworkStreamStub",
|
| 98 |
+
(),
|
| 99 |
+
{
|
| 100 |
+
"get_extra_info": lambda self, name: (
|
| 101 |
+
(peer_address, 443) if name == "server_addr" else None
|
| 102 |
+
)
|
| 103 |
+
},
|
| 104 |
+
)()
|
| 105 |
+
}
|
| 106 |
+
|
| 107 |
+
def __enter__(self) -> "_FakeStreamResponse":
|
| 108 |
+
"""进入流式响应上下文并返回自身。"""
|
| 109 |
+
return self
|
| 110 |
+
|
| 111 |
+
def __exit__(self, exc_type, exc, traceback) -> bool:
|
| 112 |
+
"""退出响应上下文且不吞掉异常。"""
|
| 113 |
+
del exc_type, exc, traceback
|
| 114 |
+
return False
|
| 115 |
+
|
| 116 |
+
def iter_bytes(self):
|
| 117 |
+
"""以单个分块返回完整响应体。"""
|
| 118 |
+
yield self.body
|
| 119 |
+
|
| 120 |
+
|
| 121 |
+
class ImageInversionJobAPITest(unittest.TestCase):
|
| 122 |
+
"""验证 Image Inversion 自定义任务 API 的公开契约与安全边界。"""
|
| 123 |
+
|
| 124 |
+
def setUp(self) -> None:
|
| 125 |
+
"""登记每个测试需要在结束时关闭的客户端和线程事件。"""
|
| 126 |
+
self.clients: list[TestClient] = []
|
| 127 |
+
self.release_events: list[threading.Event] = []
|
| 128 |
+
|
| 129 |
+
def tearDown(self) -> None:
|
| 130 |
+
"""释放测试线程和 TestClient,不创建持久化测试产物。"""
|
| 131 |
+
for release_event in self.release_events:
|
| 132 |
+
release_event.set()
|
| 133 |
+
time.sleep(0.02)
|
| 134 |
+
for client in self.clients:
|
| 135 |
+
client.close()
|
| 136 |
+
|
| 137 |
+
@staticmethod
|
| 138 |
+
def _settings(**overrides) -> ImageInversionJobSettings:
|
| 139 |
+
"""构造不依赖进程环境的隔离任务设置。
|
| 140 |
+
|
| 141 |
+
Args:
|
| 142 |
+
overrides: 要覆盖的 ImageInversionJobSettings 字段。
|
| 143 |
+
|
| 144 |
+
Returns:
|
| 145 |
+
可供单元测试直接注入的设置。
|
| 146 |
+
"""
|
| 147 |
+
values = {
|
| 148 |
+
"api_key": API_KEY,
|
| 149 |
+
"allowed_hosts": frozenset({"allowed.example"}),
|
| 150 |
+
"result_ttl_seconds": 1800,
|
| 151 |
+
"poll_after_seconds": 1,
|
| 152 |
+
"fetch_timeout_seconds": 15.0,
|
| 153 |
+
"max_image_bytes": 8 * 1024 * 1024,
|
| 154 |
+
"max_image_pixels": 4_194_304,
|
| 155 |
+
"max_image_dimension": 4096,
|
| 156 |
+
"max_records": 256,
|
| 157 |
+
"space_host": "https://space.example",
|
| 158 |
+
}
|
| 159 |
+
values.update(overrides)
|
| 160 |
+
return ImageInversionJobSettings(**values)
|
| 161 |
+
|
| 162 |
+
@staticmethod
|
| 163 |
+
def _payload(**overrides) -> dict:
|
| 164 |
+
"""生成可直接提交给创建接口的合法请求体。
|
| 165 |
+
|
| 166 |
+
Args:
|
| 167 |
+
overrides: 要覆盖或增加的 JSON 字段。
|
| 168 |
+
|
| 169 |
+
Returns:
|
| 170 |
+
包含远程输入图和全部展示参数的请求字典。
|
| 171 |
+
"""
|
| 172 |
+
payload = {
|
| 173 |
+
"input_image_url": "https://allowed.example/input.png?source=private",
|
| 174 |
+
"general_threshold": 0.35,
|
| 175 |
+
"character_threshold": 0.85,
|
| 176 |
+
"show_confidence": True,
|
| 177 |
+
"show_general": True,
|
| 178 |
+
"show_character": True,
|
| 179 |
+
"show_ip": True,
|
| 180 |
+
"separator": "comma",
|
| 181 |
+
"show_chinese": True,
|
| 182 |
+
}
|
| 183 |
+
payload.update(overrides)
|
| 184 |
+
return payload
|
| 185 |
+
|
| 186 |
+
@staticmethod
|
| 187 |
+
def _headers(**overrides) -> dict[str, str]:
|
| 188 |
+
"""生成包含共享密钥的请求头。
|
| 189 |
+
|
| 190 |
+
Args:
|
| 191 |
+
overrides: 要覆盖或追加的请求头。
|
| 192 |
+
|
| 193 |
+
Returns:
|
| 194 |
+
可传给 TestClient 的请求头字典。
|
| 195 |
+
"""
|
| 196 |
+
headers = {"X-API-Key": API_KEY}
|
| 197 |
+
headers.update(overrides)
|
| 198 |
+
return headers
|
| 199 |
+
|
| 200 |
+
def _make_service(
|
| 201 |
+
self,
|
| 202 |
+
*,
|
| 203 |
+
executor=None,
|
| 204 |
+
image_fetcher=None,
|
| 205 |
+
inference_slot: threading.Lock | None = None,
|
| 206 |
+
settings: ImageInversionJobSettings | None = None,
|
| 207 |
+
add_catch_all: bool = False,
|
| 208 |
+
) -> tuple[TestClient, ImageInversionJobAPI, list[_TrackedImage], dict[str, int]]:
|
| 209 |
+
"""创建完全隔离且不导入模型的 FastAPI 测试服务。
|
| 210 |
+
|
| 211 |
+
Args:
|
| 212 |
+
executor: 可选假标签执行器。
|
| 213 |
+
image_fetcher: 可选假远程图片抓取器。
|
| 214 |
+
inference_slot: 可选 UI/API 共享锁。
|
| 215 |
+
settings: 可选任务设置。
|
| 216 |
+
add_catch_all: 是否先注册模拟 Gradio 的宽泛路由。
|
| 217 |
+
|
| 218 |
+
Returns:
|
| 219 |
+
TestClient、API 对象、输入图片列表与调用计数器。
|
| 220 |
+
"""
|
| 221 |
+
settings = settings or self._settings()
|
| 222 |
+
created_images: list[_TrackedImage] = []
|
| 223 |
+
counters = {"fetch": 0, "execute": 0}
|
| 224 |
+
counter_lock = threading.Lock()
|
| 225 |
+
|
| 226 |
+
if image_fetcher is None:
|
| 227 |
+
|
| 228 |
+
def image_fetcher(url, current_settings):
|
| 229 |
+
"""返回可追踪关闭状态的假输入图。"""
|
| 230 |
+
del url, current_settings
|
| 231 |
+
with counter_lock:
|
| 232 |
+
counters["fetch"] += 1
|
| 233 |
+
image = _TrackedImage()
|
| 234 |
+
created_images.append(image)
|
| 235 |
+
return image
|
| 236 |
+
|
| 237 |
+
if executor is None:
|
| 238 |
+
|
| 239 |
+
def executor(payload, input_image, job_id):
|
| 240 |
+
"""返回与新接口契约一致的六字段 JSON 结果。"""
|
| 241 |
+
del payload, input_image, job_id
|
| 242 |
+
with counter_lock:
|
| 243 |
+
counters["execute"] += 1
|
| 244 |
+
return EXPECTED_RESULT
|
| 245 |
+
|
| 246 |
+
api = ImageInversionJobAPI(
|
| 247 |
+
settings=settings,
|
| 248 |
+
executor=executor,
|
| 249 |
+
inference_slot=inference_slot or threading.Lock(),
|
| 250 |
+
image_fetcher=image_fetcher,
|
| 251 |
+
)
|
| 252 |
+
app = FastAPI()
|
| 253 |
+
if add_catch_all:
|
| 254 |
+
|
| 255 |
+
@app.api_route("/{path:path}", methods=["GET", "POST"])
|
| 256 |
+
def catch_all(path: str):
|
| 257 |
+
"""模拟可能抢先匹配 API 请求的 Gradio 宽泛路由。"""
|
| 258 |
+
return {"catch_all": path}
|
| 259 |
+
|
| 260 |
+
api.install_on_app(app)
|
| 261 |
+
client = TestClient(app, base_url="https://space.example")
|
| 262 |
+
self.clients.append(client)
|
| 263 |
+
return client, api, created_images, counters
|
| 264 |
+
|
| 265 |
+
def _wait_for_terminal(
|
| 266 |
+
self,
|
| 267 |
+
client: TestClient,
|
| 268 |
+
job_id: str,
|
| 269 |
+
timeout_seconds: float = 3.0,
|
| 270 |
+
):
|
| 271 |
+
"""轮询测试任务直到返回成功或失败终态。
|
| 272 |
+
|
| 273 |
+
Args:
|
| 274 |
+
client: 当前隔离服务客户端。
|
| 275 |
+
job_id: 创建接口返回的任务标识。
|
| 276 |
+
timeout_seconds: 允许后台测试线程运行的最长时间。
|
| 277 |
+
|
| 278 |
+
Returns:
|
| 279 |
+
第一个非 202 状态响应。
|
| 280 |
+
"""
|
| 281 |
+
deadline = time.monotonic() + timeout_seconds
|
| 282 |
+
while time.monotonic() < deadline:
|
| 283 |
+
response = client.get(
|
| 284 |
+
f"/api/jobs/{job_id}",
|
| 285 |
+
headers={"X-API-Key": API_KEY},
|
| 286 |
+
)
|
| 287 |
+
if response.status_code != 202:
|
| 288 |
+
return response
|
| 289 |
+
time.sleep(0.01)
|
| 290 |
+
self.fail(f"job {job_id} did not reach a terminal state")
|
| 291 |
+
|
| 292 |
+
def test_auth_schema_cpu_contract_and_route_precedence(self) -> None:
|
| 293 |
+
"""验证鉴权、CPU 请求契约及自定义路由优先级。"""
|
| 294 |
+
client, api, _, counters = self._make_service(add_catch_all=True)
|
| 295 |
+
self.assertTrue(callable(create_job_api_lifespan(api)))
|
| 296 |
+
|
| 297 |
+
unauthorized = client.post("/api/jobs", json=self._payload())
|
| 298 |
+
self.assertEqual(unauthorized.status_code, 401)
|
| 299 |
+
self.assertEqual(counters["fetch"], 0)
|
| 300 |
+
|
| 301 |
+
unknown_field = client.post(
|
| 302 |
+
"/api/jobs",
|
| 303 |
+
json=self._payload(unknown=True),
|
| 304 |
+
headers=self._headers(),
|
| 305 |
+
)
|
| 306 |
+
self.assertEqual(unknown_field.status_code, 422)
|
| 307 |
+
self.assertEqual(counters["fetch"], 0)
|
| 308 |
+
|
| 309 |
+
# cpu-basic Space 不依赖 ZeroGPU 的 x-ip-token。
|
| 310 |
+
accepted = client.post(
|
| 311 |
+
"/api/jobs",
|
| 312 |
+
json=self._payload(),
|
| 313 |
+
headers=self._headers(),
|
| 314 |
+
)
|
| 315 |
+
self.assertEqual(accepted.status_code, 202)
|
| 316 |
+
self.assertNotIn("catch_all", accepted.json())
|
| 317 |
+
terminal = self._wait_for_terminal(client, accepted.json()["job_id"])
|
| 318 |
+
self.assertEqual(terminal.status_code, 200)
|
| 319 |
+
self.assertEqual({key: terminal.json()[key] for key in EXPECTED_RESULT}, EXPECTED_RESULT)
|
| 320 |
+
self.assertEqual(terminal.json()["status"], "succeeded")
|
| 321 |
+
self.assertEqual(client.get(terminal.json()["result_url"]).json(), EXPECTED_RESULT)
|
| 322 |
+
|
| 323 |
+
def test_success_contract_does_not_retain_caller_credentials(self) -> None:
|
| 324 |
+
"""验证成功链路仅向执行器传业务参数且结果无包装层。"""
|
| 325 |
+
seen: dict = {}
|
| 326 |
+
|
| 327 |
+
def executor(payload, input_image, job_id):
|
| 328 |
+
"""记录后台适配器参数并返回固定六字段结果。"""
|
| 329 |
+
seen.update({"payload": payload, "input_image": input_image, "job_id": job_id})
|
| 330 |
+
return EXPECTED_RESULT
|
| 331 |
+
|
| 332 |
+
client, api, images, _ = self._make_service(executor=executor)
|
| 333 |
+
created = client.post(
|
| 334 |
+
"/api/jobs",
|
| 335 |
+
json=self._payload(separator="newline"),
|
| 336 |
+
headers=self._headers(
|
| 337 |
+
Authorization="Bearer must-not-be-retained",
|
| 338 |
+
Cookie="must-not-be-retained",
|
| 339 |
+
**{"X-IP-Token": "not-needed-on-cpu"},
|
| 340 |
+
),
|
| 341 |
+
)
|
| 342 |
+
self.assertEqual(created.status_code, 202)
|
| 343 |
+
self.assertEqual(created.headers["Cache-Control"], "no-store")
|
| 344 |
+
self.assertEqual(created.headers["Location"], created.json()["status_url"])
|
| 345 |
+
job_id = created.json()["job_id"]
|
| 346 |
+
terminal = self._wait_for_terminal(client, job_id)
|
| 347 |
+
self.assertEqual(terminal.status_code, 200)
|
| 348 |
+
self.assertEqual(set(terminal.json()), set(EXPECTED_RESULT) | {"job_id", "status", "result_url"})
|
| 349 |
+
self.assertEqual(seen["job_id"], job_id)
|
| 350 |
+
self.assertEqual(seen["payload"].separator, "newline")
|
| 351 |
+
self.assertTrue(all(image.closed for image in images))
|
| 352 |
+
self.assertFalse(hasattr(api.jobs[job_id], "zero_gpu_headers"))
|
| 353 |
+
|
| 354 |
+
def test_idempotent_replay_conflict_and_busy_submission(self) -> None:
|
| 355 |
+
"""验证幂等重放、冲突及单任务繁忙状态。"""
|
| 356 |
+
release_event = threading.Event()
|
| 357 |
+
started_event = threading.Event()
|
| 358 |
+
self.release_events.append(release_event)
|
| 359 |
+
execution_count = 0
|
| 360 |
+
count_lock = threading.Lock()
|
| 361 |
+
|
| 362 |
+
def blocking_executor(payload, input_image, job_id):
|
| 363 |
+
"""阻塞首个任务,使测试可观察幂等和忙状态。"""
|
| 364 |
+
nonlocal execution_count
|
| 365 |
+
del payload, input_image, job_id
|
| 366 |
+
with count_lock:
|
| 367 |
+
execution_count += 1
|
| 368 |
+
started_event.set()
|
| 369 |
+
release_event.wait(timeout=3)
|
| 370 |
+
return EXPECTED_RESULT
|
| 371 |
+
|
| 372 |
+
client, _, _, counters = self._make_service(executor=blocking_executor)
|
| 373 |
+
headers = self._headers(**{"Idempotency-Key": "same-request"})
|
| 374 |
+
first = client.post("/api/jobs", json=self._payload(), headers=headers)
|
| 375 |
+
self.assertEqual(first.status_code, 202)
|
| 376 |
+
self.assertTrue(started_event.wait(timeout=1))
|
| 377 |
+
|
| 378 |
+
replay = client.post("/api/jobs", json=self._payload(), headers=headers)
|
| 379 |
+
self.assertEqual(replay.status_code, 202)
|
| 380 |
+
self.assertEqual(replay.json()["job_id"], first.json()["job_id"])
|
| 381 |
+
self.assertEqual(counters["fetch"], 1)
|
| 382 |
+
|
| 383 |
+
conflict = client.post(
|
| 384 |
+
"/api/jobs",
|
| 385 |
+
json=self._payload(general_threshold=0.5),
|
| 386 |
+
headers=headers,
|
| 387 |
+
)
|
| 388 |
+
self.assertEqual(conflict.status_code, 409)
|
| 389 |
+
|
| 390 |
+
busy = client.post(
|
| 391 |
+
"/api/jobs",
|
| 392 |
+
json=self._payload(character_threshold=0.7),
|
| 393 |
+
headers=self._headers(**{"Idempotency-Key": "different-request"}),
|
| 394 |
+
)
|
| 395 |
+
self.assertEqual(busy.status_code, 503)
|
| 396 |
+
self.assertEqual(busy.headers["Retry-After"], "5")
|
| 397 |
+
self.assertEqual(counters["fetch"], 1)
|
| 398 |
+
|
| 399 |
+
release_event.set()
|
| 400 |
+
terminal = self._wait_for_terminal(client, first.json()["job_id"])
|
| 401 |
+
self.assertEqual(terminal.status_code, 200)
|
| 402 |
+
self.assertEqual(execution_count, 1)
|
| 403 |
+
|
| 404 |
+
def test_concurrent_same_idempotency_key_executes_once(self) -> None:
|
| 405 |
+
"""验证并发同键请求最终只执行一次。"""
|
| 406 |
+
fetch_barrier = threading.Barrier(2)
|
| 407 |
+
fetch_count = 0
|
| 408 |
+
execute_count = 0
|
| 409 |
+
count_lock = threading.Lock()
|
| 410 |
+
|
| 411 |
+
def racing_fetcher(url, settings):
|
| 412 |
+
"""让两个请求都完成第一次幂等检查后再返回图片。"""
|
| 413 |
+
nonlocal fetch_count
|
| 414 |
+
del url, settings
|
| 415 |
+
with count_lock:
|
| 416 |
+
fetch_count += 1
|
| 417 |
+
fetch_barrier.wait(timeout=2)
|
| 418 |
+
return _TrackedImage()
|
| 419 |
+
|
| 420 |
+
def counting_executor(payload, input_image, job_id):
|
| 421 |
+
"""统计实际执行次数并返回固定结果。"""
|
| 422 |
+
nonlocal execute_count
|
| 423 |
+
del payload, input_image, job_id
|
| 424 |
+
with count_lock:
|
| 425 |
+
execute_count += 1
|
| 426 |
+
return EXPECTED_RESULT
|
| 427 |
+
|
| 428 |
+
client, _, _, _ = self._make_service(
|
| 429 |
+
image_fetcher=racing_fetcher,
|
| 430 |
+
executor=counting_executor,
|
| 431 |
+
)
|
| 432 |
+
headers = self._headers(**{"Idempotency-Key": "concurrent-key"})
|
| 433 |
+
|
| 434 |
+
def submit_job(_index: int):
|
| 435 |
+
"""提交一个并发幂等请求。"""
|
| 436 |
+
return client.post("/api/jobs", json=self._payload(), headers=headers)
|
| 437 |
+
|
| 438 |
+
with ThreadPoolExecutor(max_workers=2) as pool:
|
| 439 |
+
responses = list(pool.map(submit_job, range(2)))
|
| 440 |
+
self.assertEqual([response.status_code for response in responses], [202, 202])
|
| 441 |
+
self.assertEqual(responses[0].json()["job_id"], responses[1].json()["job_id"])
|
| 442 |
+
terminal = self._wait_for_terminal(client, responses[0].json()["job_id"])
|
| 443 |
+
self.assertEqual(terminal.status_code, 200)
|
| 444 |
+
self.assertEqual(fetch_count, 2)
|
| 445 |
+
self.assertEqual(execute_count, 1)
|
| 446 |
+
|
| 447 |
+
def test_failure_is_sanitized_and_releases_image_and_slot(self) -> None:
|
| 448 |
+
"""验证内部异常不泄密,且图片与共享锁始终释放。"""
|
| 449 |
+
shared_slot = threading.Lock()
|
| 450 |
+
|
| 451 |
+
def failing_executor(payload, input_image, job_id):
|
| 452 |
+
"""模拟包含敏感内部信息的预测错误。"""
|
| 453 |
+
del payload, input_image, job_id
|
| 454 |
+
raise RuntimeError("INTERNAL-SECRET-STACK-DATA")
|
| 455 |
+
|
| 456 |
+
client, api, images, _ = self._make_service(
|
| 457 |
+
executor=failing_executor,
|
| 458 |
+
inference_slot=shared_slot,
|
| 459 |
+
)
|
| 460 |
+
with self.assertLogs("image_inversion_job_api", level=logging.ERROR) as logs:
|
| 461 |
+
created = client.post(
|
| 462 |
+
"/api/jobs",
|
| 463 |
+
json=self._payload(),
|
| 464 |
+
headers=self._headers(Authorization="Bearer sensitive"),
|
| 465 |
+
)
|
| 466 |
+
terminal = self._wait_for_terminal(client, created.json()["job_id"])
|
| 467 |
+
|
| 468 |
+
self.assertEqual(terminal.status_code, 500)
|
| 469 |
+
self.assertEqual(terminal.json(), {"job_id": created.json()["job_id"], "status": "failed", "error": {"code": "GENERATION_FAILED"}})
|
| 470 |
+
combined_logs = "\n".join(logs.output)
|
| 471 |
+
self.assertNotIn("INTERNAL-SECRET-STACK-DATA", combined_logs)
|
| 472 |
+
self.assertNotIn(API_KEY, combined_logs)
|
| 473 |
+
self.assertTrue(all(image.closed for image in images))
|
| 474 |
+
self.assertTrue(shared_slot.acquire(blocking=False))
|
| 475 |
+
shared_slot.release()
|
| 476 |
+
self.assertEqual(api.jobs[created.json()["job_id"]].error_type, "RuntimeError")
|
| 477 |
+
|
| 478 |
+
def test_thread_start_failure_rolls_back_and_retry_succeeds(self) -> None:
|
| 479 |
+
"""验证后台线程启动失败时任务、幂等索引和锁均回滚。"""
|
| 480 |
+
client, api, images, counters = self._make_service()
|
| 481 |
+
with patch(
|
| 482 |
+
"image_inversion_job_api.threading.Thread.start",
|
| 483 |
+
side_effect=RuntimeError("thread-start-failed"),
|
| 484 |
+
):
|
| 485 |
+
with self.assertRaises(RuntimeError):
|
| 486 |
+
client.post(
|
| 487 |
+
"/api/jobs",
|
| 488 |
+
json=self._payload(),
|
| 489 |
+
headers=self._headers(**{"Idempotency-Key": "rollback-key"}),
|
| 490 |
+
)
|
| 491 |
+
|
| 492 |
+
self.assertEqual(api.jobs, {})
|
| 493 |
+
self.assertEqual(api.idempotency_jobs, {})
|
| 494 |
+
self.assertTrue(all(image.closed for image in images))
|
| 495 |
+
self.assertFalse(api.inference_slot.locked())
|
| 496 |
+
|
| 497 |
+
retry = client.post(
|
| 498 |
+
"/api/jobs",
|
| 499 |
+
json=self._payload(),
|
| 500 |
+
headers=self._headers(**{"Idempotency-Key": "rollback-key"}),
|
| 501 |
+
)
|
| 502 |
+
self.assertEqual(retry.status_code, 202)
|
| 503 |
+
self.assertEqual(
|
| 504 |
+
self._wait_for_terminal(client, retry.json()["job_id"]).status_code,
|
| 505 |
+
200,
|
| 506 |
+
)
|
| 507 |
+
self.assertEqual(counters["execute"], 1)
|
| 508 |
+
|
| 509 |
+
def test_ttl_removes_all_expired_terminal_records(self) -> None:
|
| 510 |
+
"""验证 TTL 清理会移除全部已过期终态记录。"""
|
| 511 |
+
settings = self._settings(result_ttl_seconds=60)
|
| 512 |
+
client, api, _, _ = self._make_service(settings=settings)
|
| 513 |
+
job_ids: list[str] = []
|
| 514 |
+
for threshold in (0.2, 0.3):
|
| 515 |
+
created = client.post(
|
| 516 |
+
"/api/jobs",
|
| 517 |
+
json=self._payload(general_threshold=threshold),
|
| 518 |
+
headers=self._headers(),
|
| 519 |
+
)
|
| 520 |
+
job_id = created.json()["job_id"]
|
| 521 |
+
self.assertEqual(self._wait_for_terminal(client, job_id).status_code, 200)
|
| 522 |
+
api.jobs[job_id].completed_at = time.time() - 100
|
| 523 |
+
job_ids.append(job_id)
|
| 524 |
+
|
| 525 |
+
api.remove_expired_jobs()
|
| 526 |
+
self.assertEqual(api.jobs, {})
|
| 527 |
+
self.assertEqual(len(job_ids), 2)
|
| 528 |
+
|
| 529 |
+
def test_capacity_evicts_oldest_terminal_and_preserves_running_records(self) -> None:
|
| 530 |
+
"""验证容量上限只淘汰最老终态且满是运行态时返回 503。"""
|
| 531 |
+
settings = self._settings(max_records=2)
|
| 532 |
+
client, api, _, _ = self._make_service(settings=settings)
|
| 533 |
+
completed_job_ids: list[str] = []
|
| 534 |
+
for threshold in (0.2, 0.3):
|
| 535 |
+
created = client.post(
|
| 536 |
+
"/api/jobs",
|
| 537 |
+
json=self._payload(general_threshold=threshold),
|
| 538 |
+
headers=self._headers(),
|
| 539 |
+
)
|
| 540 |
+
job_id = created.json()["job_id"]
|
| 541 |
+
self.assertEqual(self._wait_for_terminal(client, job_id).status_code, 200)
|
| 542 |
+
completed_job_ids.append(job_id)
|
| 543 |
+
time.sleep(0.01)
|
| 544 |
+
|
| 545 |
+
third = client.post(
|
| 546 |
+
"/api/jobs",
|
| 547 |
+
json=self._payload(general_threshold=0.4),
|
| 548 |
+
headers=self._headers(),
|
| 549 |
+
)
|
| 550 |
+
self.assertEqual(third.status_code, 202)
|
| 551 |
+
self.assertNotIn(completed_job_ids[0], api.jobs)
|
| 552 |
+
self.assertIn(completed_job_ids[1], api.jobs)
|
| 553 |
+
self.assertLessEqual(len(api.jobs), 2)
|
| 554 |
+
self.assertEqual(self._wait_for_terminal(client, third.json()["job_id"]).status_code, 200)
|
| 555 |
+
|
| 556 |
+
full_settings = self._settings(max_records=1)
|
| 557 |
+
full_client, full_api, _, counters = self._make_service(settings=full_settings)
|
| 558 |
+
running_record = ImageInversionJobRecord(
|
| 559 |
+
job_id="running-record",
|
| 560 |
+
payload=ImageInversionJobRequest(**self._payload()),
|
| 561 |
+
request_fingerprint="fingerprint",
|
| 562 |
+
input_image=_TrackedImage(),
|
| 563 |
+
status="running",
|
| 564 |
+
)
|
| 565 |
+
full_api.jobs[running_record.job_id] = running_record
|
| 566 |
+
blocked = full_client.post(
|
| 567 |
+
"/api/jobs",
|
| 568 |
+
json=self._payload(general_threshold=0.6),
|
| 569 |
+
headers=self._headers(),
|
| 570 |
+
)
|
| 571 |
+
self.assertEqual(blocked.status_code, 503)
|
| 572 |
+
self.assertEqual(blocked.headers["Retry-After"], "5")
|
| 573 |
+
self.assertIn(running_record.job_id, full_api.jobs)
|
| 574 |
+
self.assertEqual(counters["fetch"], 0)
|
| 575 |
+
|
| 576 |
+
def test_settings_schema_and_fingerprint(self) -> None:
|
| 577 |
+
"""验证环境设置、公开 Schema 和请求指纹边界。"""
|
| 578 |
+
with patch.dict(
|
| 579 |
+
os.environ,
|
| 580 |
+
{
|
| 581 |
+
"JOB_API_KEY": API_KEY,
|
| 582 |
+
"JOB_IMAGE_ALLOWED_HOSTS": "Allowed.Example,cdn.example",
|
| 583 |
+
},
|
| 584 |
+
clear=True,
|
| 585 |
+
):
|
| 586 |
+
settings = ImageInversionJobSettings.from_env()
|
| 587 |
+
self.assertEqual(
|
| 588 |
+
settings.allowed_hosts,
|
| 589 |
+
frozenset({"allowed.example", "cdn.example"}),
|
| 590 |
+
)
|
| 591 |
+
self.assertEqual(settings.max_image_bytes, 8 * 1024 * 1024)
|
| 592 |
+
self.assertEqual(settings.max_records, 256)
|
| 593 |
+
|
| 594 |
+
for invalid_max_records in (7, 4097):
|
| 595 |
+
with self.subTest(max_records=invalid_max_records), patch.dict(
|
| 596 |
+
os.environ,
|
| 597 |
+
{
|
| 598 |
+
"JOB_API_KEY": API_KEY,
|
| 599 |
+
"JOB_IMAGE_ALLOWED_HOSTS": "allowed.example",
|
| 600 |
+
"JOB_MAX_RECORDS": str(invalid_max_records),
|
| 601 |
+
},
|
| 602 |
+
clear=True,
|
| 603 |
+
):
|
| 604 |
+
with self.assertRaises(RuntimeError):
|
| 605 |
+
ImageInversionJobSettings.from_env()
|
| 606 |
+
|
| 607 |
+
minimum = ImageInversionJobRequest(
|
| 608 |
+
input_image_url="https://allowed.example/input.png",
|
| 609 |
+
general_threshold=0.0,
|
| 610 |
+
character_threshold=0.0,
|
| 611 |
+
separator="space",
|
| 612 |
+
)
|
| 613 |
+
maximum = ImageInversionJobRequest(
|
| 614 |
+
input_image_url="https://allowed.example/input.png",
|
| 615 |
+
general_threshold=1.0,
|
| 616 |
+
character_threshold=1.0,
|
| 617 |
+
separator="newline",
|
| 618 |
+
)
|
| 619 |
+
self.assertNotEqual(fingerprint_request(minimum), fingerprint_request(maximum))
|
| 620 |
+
reordered = ImageInversionJobRequest.model_validate(
|
| 621 |
+
minimum.model_dump(mode="json")
|
| 622 |
+
)
|
| 623 |
+
self.assertEqual(fingerprint_request(minimum), fingerprint_request(reordered))
|
| 624 |
+
|
| 625 |
+
for invalid_payload in (
|
| 626 |
+
self._payload(input_image_url="http://allowed.example/input.png"),
|
| 627 |
+
self._payload(general_threshold=-0.01),
|
| 628 |
+
self._payload(character_threshold=1.01),
|
| 629 |
+
self._payload(general_threshold=float("nan")),
|
| 630 |
+
self._payload(separator="逗号"),
|
| 631 |
+
self._payload(unknown=True),
|
| 632 |
+
):
|
| 633 |
+
with self.subTest(payload=invalid_payload):
|
| 634 |
+
with self.assertRaises(ValidationError):
|
| 635 |
+
ImageInversionJobRequest(**invalid_payload)
|
| 636 |
+
|
| 637 |
+
def test_url_policy_decode_and_fetch_do_not_leak_headers(self) -> None:
|
| 638 |
+
"""验证白名单、SSRF、图片解码和无凭据抓取边界。"""
|
| 639 |
+
settings = self._settings()
|
| 640 |
+
public_record = (
|
| 641 |
+
socket.AF_INET,
|
| 642 |
+
socket.SOCK_STREAM,
|
| 643 |
+
socket.IPPROTO_TCP,
|
| 644 |
+
"",
|
| 645 |
+
("93.184.216.34", 443),
|
| 646 |
+
)
|
| 647 |
+
with patch(
|
| 648 |
+
"image_inversion_job_api.socket.getaddrinfo",
|
| 649 |
+
return_value=[public_record],
|
| 650 |
+
):
|
| 651 |
+
validated = validate_image_url(
|
| 652 |
+
"https://allowed.example/image.png?private=query",
|
| 653 |
+
settings,
|
| 654 |
+
)
|
| 655 |
+
self.assertEqual(validated.host, "allowed.example")
|
| 656 |
+
for invalid_url in (
|
| 657 |
+
"http://allowed.example/image.png",
|
| 658 |
+
"https://user:password@allowed.example/image.png",
|
| 659 |
+
"https://allowed.example:444/image.png",
|
| 660 |
+
"https://sub.allowed.example/image.png",
|
| 661 |
+
"https://allowed.example/image.png#fragment",
|
| 662 |
+
):
|
| 663 |
+
with self.subTest(url=invalid_url):
|
| 664 |
+
with self.assertRaises(ImageSourceError):
|
| 665 |
+
validate_image_url(invalid_url, settings)
|
| 666 |
+
|
| 667 |
+
private_record = (
|
| 668 |
+
socket.AF_INET,
|
| 669 |
+
socket.SOCK_STREAM,
|
| 670 |
+
socket.IPPROTO_TCP,
|
| 671 |
+
"",
|
| 672 |
+
("127.0.0.1", 443),
|
| 673 |
+
)
|
| 674 |
+
with patch(
|
| 675 |
+
"image_inversion_job_api.socket.getaddrinfo",
|
| 676 |
+
return_value=[public_record, private_record],
|
| 677 |
+
):
|
| 678 |
+
with self.assertRaises(ImageSourceError):
|
| 679 |
+
validate_image_url("https://allowed.example/image.png", settings)
|
| 680 |
+
|
| 681 |
+
decoded = decode_image_bytes(PNG_BYTES, "application/octet-stream", settings)
|
| 682 |
+
self.assertEqual(decoded.mode, "RGB")
|
| 683 |
+
decoded.close()
|
| 684 |
+
|
| 685 |
+
captured: dict = {}
|
| 686 |
+
|
| 687 |
+
class FakeClient:
|
| 688 |
+
"""捕获 httpx 客户端配置并返回假响应。"""
|
| 689 |
+
|
| 690 |
+
def __init__(self, **kwargs) -> None:
|
| 691 |
+
captured["client_kwargs"] = kwargs
|
| 692 |
+
|
| 693 |
+
def __enter__(self):
|
| 694 |
+
"""进入客户端上下文并返回自身。"""
|
| 695 |
+
return self
|
| 696 |
+
|
| 697 |
+
def __exit__(self, exc_type, exc, traceback) -> bool:
|
| 698 |
+
"""退出客户端上下文且不吞掉异常。"""
|
| 699 |
+
del exc_type, exc, traceback
|
| 700 |
+
return False
|
| 701 |
+
|
| 702 |
+
def close(self) -> None:
|
| 703 |
+
"""模拟关闭客户端。"""
|
| 704 |
+
captured["closed"] = True
|
| 705 |
+
|
| 706 |
+
def stream(self, method, url, headers):
|
| 707 |
+
"""记录公开请求参数并返回假响应。"""
|
| 708 |
+
captured["method"] = method
|
| 709 |
+
captured["url"] = url
|
| 710 |
+
captured["headers"] = dict(headers)
|
| 711 |
+
return captured["response"]
|
| 712 |
+
|
| 713 |
+
with patch(
|
| 714 |
+
"image_inversion_job_api.socket.getaddrinfo",
|
| 715 |
+
return_value=[public_record],
|
| 716 |
+
), patch("image_inversion_job_api.httpx.Client", FakeClient):
|
| 717 |
+
captured["response"] = _FakeStreamResponse(200, PNG_BYTES)
|
| 718 |
+
fetched = fetch_remote_image(
|
| 719 |
+
"https://allowed.example/image.png?secret=value",
|
| 720 |
+
settings,
|
| 721 |
+
)
|
| 722 |
+
fetched.close()
|
| 723 |
+
self.assertFalse(captured["client_kwargs"]["follow_redirects"])
|
| 724 |
+
self.assertFalse(captured["client_kwargs"]["trust_env"])
|
| 725 |
+
self.assertEqual(
|
| 726 |
+
captured["headers"],
|
| 727 |
+
{"Accept": "image/png,image/jpeg,image/webp"},
|
| 728 |
+
)
|
| 729 |
+
for header_name in ("Authorization", "Cookie", "X-API-Key", "X-IP-Token"):
|
| 730 |
+
self.assertNotIn(header_name, captured["headers"])
|
| 731 |
+
|
| 732 |
+
captured["response"] = _FakeStreamResponse(302, b"")
|
| 733 |
+
with self.assertRaises(ImageSourceError):
|
| 734 |
+
fetch_remote_image("https://allowed.example/redirect", settings)
|
| 735 |
+
|
| 736 |
+
captured["response"] = _FakeStreamResponse(
|
| 737 |
+
200,
|
| 738 |
+
PNG_BYTES,
|
| 739 |
+
peer_address="127.0.0.1",
|
| 740 |
+
)
|
| 741 |
+
with self.assertRaises(ImageSourceError):
|
| 742 |
+
fetch_remote_image("https://allowed.example/rebound", settings)
|
| 743 |
+
|
| 744 |
+
def test_app_wires_job_lifespan_and_scopes_hf_token(self) -> None:
|
| 745 |
+
"""静态验证应用接线、六字段结果与受限模型下载令牌边界。"""
|
| 746 |
+
app_path = Path(__file__).resolve().parents[1] / "app.py"
|
| 747 |
+
app_source = app_path.read_text(encoding="utf-8")
|
| 748 |
+
app_tree = ast.parse(app_source)
|
| 749 |
+
|
| 750 |
+
executor_function = next(
|
| 751 |
+
node
|
| 752 |
+
for node in app_tree.body
|
| 753 |
+
if isinstance(node, ast.FunctionDef)
|
| 754 |
+
and node.name == "execute_image_inversion_job"
|
| 755 |
+
)
|
| 756 |
+
result_return = next(
|
| 757 |
+
node
|
| 758 |
+
for node in ast.walk(executor_function)
|
| 759 |
+
if isinstance(node, ast.Return) and isinstance(node.value, ast.Dict)
|
| 760 |
+
)
|
| 761 |
+
result_keys = {
|
| 762 |
+
key.value
|
| 763 |
+
for key in result_return.value.keys
|
| 764 |
+
if isinstance(key, ast.Constant) and isinstance(key.value, str)
|
| 765 |
+
}
|
| 766 |
+
self.assertEqual(result_keys, set(EXPECTED_RESULT))
|
| 767 |
+
self.assertIn("ssr_mode=False", app_source)
|
| 768 |
+
self.assertIn('app_kwargs={"lifespan": JOB_API_LIFESPAN}', app_source)
|
| 769 |
+
self.assertIn('ASSETS_HF_TOKEN = os.environ.get("HF_TOKEN")', app_source)
|
| 770 |
+
self.assertIn('token=ASSETS_HF_TOKEN or False', app_source)
|
| 771 |
+
self.assertIn('HF_HUB_DISABLE_IMPLICIT_TOKEN', app_source)
|
| 772 |
+
|
| 773 |
+
|
| 774 |
+
if __name__ == "__main__":
|
| 775 |
+
unittest.main()
|