IdlecloudX commited on
Commit
b91d49e
·
verified ·
1 Parent(s): 2109375

Keep HEAD metadata consistent and update existing API contract tests

Browse files
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
- def test_success_contract_does_not_retain_caller_credentials(self) -> None:
322
- """验证成功链路仅向执行器传业务参数且结果无包装层。"""
323
- seen: dict = {}
324
-
325
- def executor(payload, input_image, job_id):
326
- """记录后台适配器参数并返回固定六字段结果。"""
327
- seen.update({"payload": payload, "input_image": input_image, "job_id": job_id})
328
- return EXPECTED_RESULT
329
-
330
- client, api, images, _ = self._make_service(executor=executor)
331
- created = client.post(
332
- "/api/jobs",
333
- json=self._payload(separator="newline"),
334
- headers=self._headers(
335
- Authorization="Bearer must-not-be-retained",
336
- Cookie="must-not-be-retained",
337
- **{"X-IP-Token": "not-needed-on-cpu"},
338
- ),
339
- )
340
- self.assertEqual(created.status_code, 202)
341
- self.assertEqual(created.headers["Cache-Control"], "no-store")
342
- self.assertEqual(created.headers["Location"], created.json()["status_url"])
343
- job_id = created.json()["job_id"]
344
- terminal = self._wait_for_terminal(client, job_id)
345
- self.assertEqual(terminal.status_code, 200)
346
- self.assertEqual(set(terminal.json()), set(EXPECTED_RESULT))
347
- self.assertEqual(seen["job_id"], job_id)
348
- self.assertEqual(seen["payload"].separator, "newline")
349
- self.assertTrue(all(image.closed for image in images))
350
- self.assertFalse(hasattr(api.jobs[job_id], "zero_gpu_headers"))
351
-
352
- def test_idempotent_replay_conflict_and_busy_submission(self) -> None:
353
- """验证幂等重放、冲突及单任务繁忙状态。"""
354
- release_event = threading.Event()
355
- started_event = threading.Event()
356
- self.release_events.append(release_event)
357
- execution_count = 0
358
- count_lock = threading.Lock()
359
-
360
- def blocking_executor(payload, input_image, job_id):
361
- """阻塞首个任务,使测试可观察幂等和忙状态。"""
362
- nonlocal execution_count
363
- del payload, input_image, job_id
364
- with count_lock:
365
- execution_count += 1
366
- started_event.set()
367
- release_event.wait(timeout=3)
368
- return EXPECTED_RESULT
369
-
370
- client, _, _, counters = self._make_service(executor=blocking_executor)
371
- headers = self._headers(**{"Idempotency-Key": "same-request"})
372
- first = client.post("/api/jobs", json=self._payload(), headers=headers)
373
- self.assertEqual(first.status_code, 202)
374
- self.assertTrue(started_event.wait(timeout=1))
375
-
376
- replay = client.post("/api/jobs", json=self._payload(), headers=headers)
377
- self.assertEqual(replay.status_code, 202)
378
- self.assertEqual(replay.json()["job_id"], first.json()["job_id"])
379
- self.assertEqual(counters["fetch"], 1)
380
-
381
- conflict = client.post(
382
- "/api/jobs",
383
- json=self._payload(general_threshold=0.5),
384
- headers=headers,
385
- )
386
- self.assertEqual(conflict.status_code, 409)
387
-
388
- busy = client.post(
389
- "/api/jobs",
390
- json=self._payload(character_threshold=0.7),
391
- headers=self._headers(**{"Idempotency-Key": "different-request"}),
392
- )
393
- self.assertEqual(busy.status_code, 503)
394
- self.assertEqual(busy.headers["Retry-After"], "5")
395
- self.assertEqual(counters["fetch"], 1)
396
-
397
- release_event.set()
398
- terminal = self._wait_for_terminal(client, first.json()["job_id"])
399
- self.assertEqual(terminal.status_code, 200)
400
- self.assertEqual(execution_count, 1)
401
-
402
- def test_concurrent_same_idempotency_key_executes_once(self) -> None:
403
- """验证并发同键请求最终只执行一次。"""
404
- fetch_barrier = threading.Barrier(2)
405
- fetch_count = 0
406
- execute_count = 0
407
- count_lock = threading.Lock()
408
-
409
- def racing_fetcher(url, settings):
410
- """让两个请求都完成第一次幂等检查后再返回图片。"""
411
- nonlocal fetch_count
412
- del url, settings
413
- with count_lock:
414
- fetch_count += 1
415
- fetch_barrier.wait(timeout=2)
416
- return _TrackedImage()
417
-
418
- def counting_executor(payload, input_image, job_id):
419
- """统计实际执行次数并返回固定结果。"""
420
- nonlocal execute_count
421
- del payload, input_image, job_id
422
- with count_lock:
423
- execute_count += 1
424
- return EXPECTED_RESULT
425
-
426
- client, _, _, _ = self._make_service(
427
- image_fetcher=racing_fetcher,
428
- executor=counting_executor,
429
- )
430
- headers = self._headers(**{"Idempotency-Key": "concurrent-key"})
431
-
432
- def submit_job(_index: int):
433
- """提交一个并发幂等请求。"""
434
- return client.post("/api/jobs", json=self._payload(), headers=headers)
435
-
436
- with ThreadPoolExecutor(max_workers=2) as pool:
437
- responses = list(pool.map(submit_job, range(2)))
438
- self.assertEqual([response.status_code for response in responses], [202, 202])
439
- self.assertEqual(responses[0].json()["job_id"], responses[1].json()["job_id"])
440
- terminal = self._wait_for_terminal(client, responses[0].json()["job_id"])
441
- self.assertEqual(terminal.status_code, 200)
442
- self.assertEqual(fetch_count, 2)
443
- self.assertEqual(execute_count, 1)
444
-
445
- def test_failure_is_sanitized_and_releases_image_and_slot(self) -> None:
446
- """验证内部异常不泄密,且图片与共享锁始终释放。"""
447
- shared_slot = threading.Lock()
448
-
449
- def failing_executor(payload, input_image, job_id):
450
- """模拟包含敏感内部信息的预测错误。"""
451
- del payload, input_image, job_id
452
- raise RuntimeError("INTERNAL-SECRET-STACK-DATA")
453
-
454
- client, api, images, _ = self._make_service(
455
- executor=failing_executor,
456
- inference_slot=shared_slot,
457
- )
458
- with self.assertLogs("image_inversion_job_api", level=logging.ERROR) as logs:
459
- created = client.post(
460
- "/api/jobs",
461
- json=self._payload(),
462
- headers=self._headers(Authorization="Bearer sensitive"),
463
- )
464
- terminal = self._wait_for_terminal(client, created.json()["job_id"])
465
-
466
- self.assertEqual(terminal.status_code, 500)
467
- self.assertEqual(terminal.json(), {"error": {"code": "PREDICTION_FAILED"}})
468
- combined_logs = "\n".join(logs.output)
469
- self.assertNotIn("INTERNAL-SECRET-STACK-DATA", combined_logs)
470
- self.assertNotIn(API_KEY, combined_logs)
471
- self.assertTrue(all(image.closed for image in images))
472
- self.assertTrue(shared_slot.acquire(blocking=False))
473
- shared_slot.release()
474
- self.assertEqual(api.jobs[created.json()["job_id"]].error_type, "RuntimeError")
475
-
476
- def test_thread_start_failure_rolls_back_and_retry_succeeds(self) -> None:
477
- """验证后台线程启动失败时任务、幂等索引和锁均回滚。"""
478
- client, api, images, counters = self._make_service()
479
- with patch(
480
- "image_inversion_job_api.threading.Thread.start",
481
- side_effect=RuntimeError("thread-start-failed"),
482
- ):
483
- with self.assertRaises(RuntimeError):
484
- client.post(
485
- "/api/jobs",
486
- json=self._payload(),
487
- headers=self._headers(**{"Idempotency-Key": "rollback-key"}),
488
- )
489
-
490
- self.assertEqual(api.jobs, {})
491
- self.assertEqual(api.idempotency_jobs, {})
492
- self.assertTrue(all(image.closed for image in images))
493
- self.assertFalse(api.inference_slot.locked())
494
-
495
- retry = client.post(
496
- "/api/jobs",
497
- json=self._payload(),
498
- headers=self._headers(**{"Idempotency-Key": "rollback-key"}),
499
- )
500
- self.assertEqual(retry.status_code, 202)
501
- self.assertEqual(
502
- self._wait_for_terminal(client, retry.json()["job_id"]).status_code,
503
- 200,
504
- )
505
- self.assertEqual(counters["execute"], 1)
506
-
507
- def test_ttl_removes_all_expired_terminal_records(self) -> None:
508
- """验证 TTL 清理会移除全部已过期终态记录。"""
509
- settings = self._settings(result_ttl_seconds=60)
510
- client, api, _, _ = self._make_service(settings=settings)
511
- job_ids: list[str] = []
512
- for threshold in (0.2, 0.3):
513
- created = client.post(
514
- "/api/jobs",
515
- json=self._payload(general_threshold=threshold),
516
- headers=self._headers(),
517
- )
518
- job_id = created.json()["job_id"]
519
- self.assertEqual(self._wait_for_terminal(client, job_id).status_code, 200)
520
- api.jobs[job_id].completed_at = time.time() - 100
521
- job_ids.append(job_id)
522
-
523
- api.remove_expired_jobs()
524
- self.assertEqual(api.jobs, {})
525
- self.assertEqual(len(job_ids), 2)
526
-
527
- def test_capacity_evicts_oldest_terminal_and_preserves_running_records(self) -> None:
528
- """验证容量上限只淘汰最老终态且满是运行态时返回 503。"""
529
- settings = self._settings(max_records=2)
530
- client, api, _, _ = self._make_service(settings=settings)
531
- completed_job_ids: list[str] = []
532
- for threshold in (0.2, 0.3):
533
- created = client.post(
534
- "/api/jobs",
535
- json=self._payload(general_threshold=threshold),
536
- headers=self._headers(),
537
- )
538
- job_id = created.json()["job_id"]
539
- self.assertEqual(self._wait_for_terminal(client, job_id).status_code, 200)
540
- completed_job_ids.append(job_id)
541
- time.sleep(0.01)
542
-
543
- third = client.post(
544
- "/api/jobs",
545
- json=self._payload(general_threshold=0.4),
546
- headers=self._headers(),
547
- )
548
- self.assertEqual(third.status_code, 202)
549
- self.assertNotIn(completed_job_ids[0], api.jobs)
550
- self.assertIn(completed_job_ids[1], api.jobs)
551
- self.assertLessEqual(len(api.jobs), 2)
552
- self.assertEqual(self._wait_for_terminal(client, third.json()["job_id"]).status_code, 200)
553
-
554
- full_settings = self._settings(max_records=1)
555
- full_client, full_api, _, counters = self._make_service(settings=full_settings)
556
- running_record = ImageInversionJobRecord(
557
- job_id="running-record",
558
- payload=ImageInversionJobRequest(**self._payload()),
559
- request_fingerprint="fingerprint",
560
- input_image=_TrackedImage(),
561
- status="running",
562
- )
563
- full_api.jobs[running_record.job_id] = running_record
564
- blocked = full_client.post(
565
- "/api/jobs",
566
- json=self._payload(general_threshold=0.6),
567
- headers=self._headers(),
568
- )
569
- self.assertEqual(blocked.status_code, 503)
570
- self.assertEqual(blocked.headers["Retry-After"], "5")
571
- self.assertIn(running_record.job_id, full_api.jobs)
572
- self.assertEqual(counters["fetch"], 0)
573
-
574
- def test_settings_schema_and_fingerprint(self) -> None:
575
- """验证环境设置、公开 Schema 和请求指纹边界。"""
576
- with patch.dict(
577
- os.environ,
578
- {
579
- "JOB_API_KEY": API_KEY,
580
- "JOB_IMAGE_ALLOWED_HOSTS": "Allowed.Example,cdn.example",
581
- },
582
- clear=True,
583
- ):
584
- settings = ImageInversionJobSettings.from_env()
585
- self.assertEqual(
586
- settings.allowed_hosts,
587
- frozenset({"allowed.example", "cdn.example"}),
588
- )
589
- self.assertEqual(settings.max_image_bytes, 8 * 1024 * 1024)
590
- self.assertEqual(settings.max_records, 256)
591
-
592
- for invalid_max_records in (7, 4097):
593
- with self.subTest(max_records=invalid_max_records), patch.dict(
594
- os.environ,
595
- {
596
- "JOB_API_KEY": API_KEY,
597
- "JOB_IMAGE_ALLOWED_HOSTS": "allowed.example",
598
- "JOB_MAX_RECORDS": str(invalid_max_records),
599
- },
600
- clear=True,
601
- ):
602
- with self.assertRaises(RuntimeError):
603
- ImageInversionJobSettings.from_env()
604
-
605
- minimum = ImageInversionJobRequest(
606
- input_image_url="https://allowed.example/input.png",
607
- general_threshold=0.0,
608
- character_threshold=0.0,
609
- separator="space",
610
- )
611
- maximum = ImageInversionJobRequest(
612
- input_image_url="https://allowed.example/input.png",
613
- general_threshold=1.0,
614
- character_threshold=1.0,
615
- separator="newline",
616
- )
617
- self.assertNotEqual(fingerprint_request(minimum), fingerprint_request(maximum))
618
- reordered = ImageInversionJobRequest.model_validate(
619
- minimum.model_dump(mode="json")
620
- )
621
- self.assertEqual(fingerprint_request(minimum), fingerprint_request(reordered))
622
-
623
- for invalid_payload in (
624
- self._payload(input_image_url="http://allowed.example/input.png"),
625
- self._payload(general_threshold=-0.01),
626
- self._payload(character_threshold=1.01),
627
- self._payload(general_threshold=float("nan")),
628
- self._payload(separator="逗号"),
629
- self._payload(unknown=True),
630
- ):
631
- with self.subTest(payload=invalid_payload):
632
- with self.assertRaises(ValidationError):
633
- ImageInversionJobRequest(**invalid_payload)
634
-
635
- def test_url_policy_decode_and_fetch_do_not_leak_headers(self) -> None:
636
- """验证白名单、SSRF、图片解码和无凭据抓取边界。"""
637
- settings = self._settings()
638
- public_record = (
639
- socket.AF_INET,
640
- socket.SOCK_STREAM,
641
- socket.IPPROTO_TCP,
642
- "",
643
- ("93.184.216.34", 443),
644
- )
645
- with patch(
646
- "image_inversion_job_api.socket.getaddrinfo",
647
- return_value=[public_record],
648
- ):
649
- validated = validate_image_url(
650
- "https://allowed.example/image.png?private=query",
651
- settings,
652
- )
653
- self.assertEqual(validated.host, "allowed.example")
654
- for invalid_url in (
655
- "http://allowed.example/image.png",
656
- "https://user:password@allowed.example/image.png",
657
- "https://allowed.example:444/image.png",
658
- "https://sub.allowed.example/image.png",
659
- "https://allowed.example/image.png#fragment",
660
- ):
661
- with self.subTest(url=invalid_url):
662
- with self.assertRaises(ImageSourceError):
663
- validate_image_url(invalid_url, settings)
664
-
665
- private_record = (
666
- socket.AF_INET,
667
- socket.SOCK_STREAM,
668
- socket.IPPROTO_TCP,
669
- "",
670
- ("127.0.0.1", 443),
671
- )
672
- with patch(
673
- "image_inversion_job_api.socket.getaddrinfo",
674
- return_value=[public_record, private_record],
675
- ):
676
- with self.assertRaises(ImageSourceError):
677
- validate_image_url("https://allowed.example/image.png", settings)
678
-
679
- decoded = decode_image_bytes(PNG_BYTES, "application/octet-stream", settings)
680
- self.assertEqual(decoded.mode, "RGB")
681
- decoded.close()
682
-
683
- captured: dict = {}
684
-
685
- class FakeClient:
686
- """捕获 httpx 客户端配置并返回假响应。"""
687
-
688
- def __init__(self, **kwargs) -> None:
689
- captured["client_kwargs"] = kwargs
690
-
691
- def __enter__(self):
692
- """进入客户端上下文并返回自身。"""
693
- return self
694
-
695
- def __exit__(self, exc_type, exc, traceback) -> bool:
696
- """退出客户端上下文且不吞掉异常。"""
697
- del exc_type, exc, traceback
698
- return False
699
-
700
- def close(self) -> None:
701
- """模拟关闭客户端。"""
702
- captured["closed"] = True
703
-
704
- def stream(self, method, url, headers):
705
- """记录公开请求参数并返回假响应。"""
706
- captured["method"] = method
707
- captured["url"] = url
708
- captured["headers"] = dict(headers)
709
- return captured["response"]
710
-
711
- with patch(
712
- "image_inversion_job_api.socket.getaddrinfo",
713
- return_value=[public_record],
714
- ), patch("image_inversion_job_api.httpx.Client", FakeClient):
715
- captured["response"] = _FakeStreamResponse(200, PNG_BYTES)
716
- fetched = fetch_remote_image(
717
- "https://allowed.example/image.png?secret=value",
718
- settings,
719
- )
720
- fetched.close()
721
- self.assertFalse(captured["client_kwargs"]["follow_redirects"])
722
- self.assertFalse(captured["client_kwargs"]["trust_env"])
723
- self.assertEqual(
724
- captured["headers"],
725
- {"Accept": "image/png,image/jpeg,image/webp"},
726
- )
727
- for header_name in ("Authorization", "Cookie", "X-API-Key", "X-IP-Token"):
728
- self.assertNotIn(header_name, captured["headers"])
729
-
730
- captured["response"] = _FakeStreamResponse(302, b"")
731
- with self.assertRaises(ImageSourceError):
732
- fetch_remote_image("https://allowed.example/redirect", settings)
733
-
734
- captured["response"] = _FakeStreamResponse(
735
- 200,
736
- PNG_BYTES,
737
- peer_address="127.0.0.1",
738
- )
739
- with self.assertRaises(ImageSourceError):
740
- fetch_remote_image("https://allowed.example/rebound", settings)
741
-
742
- def test_app_wires_job_lifespan_and_scopes_hf_token(self) -> None:
743
- """静态验证应用接线、六字段结果与受限模型下载令牌边界。"""
744
- app_path = Path(__file__).resolve().parents[1] / "app.py"
745
- app_source = app_path.read_text(encoding="utf-8")
746
- app_tree = ast.parse(app_source)
747
-
748
- executor_function = next(
749
- node
750
- for node in app_tree.body
751
- if isinstance(node, ast.FunctionDef)
752
- and node.name == "execute_image_inversion_job"
753
- )
754
- result_return = next(
755
- node
756
- for node in ast.walk(executor_function)
757
- if isinstance(node, ast.Return) and isinstance(node.value, ast.Dict)
758
- )
759
- result_keys = {
760
- key.value
761
- for key in result_return.value.keys
762
- if isinstance(key, ast.Constant) and isinstance(key.value, str)
763
- }
764
- self.assertEqual(result_keys, set(EXPECTED_RESULT))
765
- self.assertIn("ssr_mode=False", app_source)
766
- self.assertIn('app_kwargs={"lifespan": JOB_API_LIFESPAN}', app_source)
767
- self.assertIn('ASSETS_HF_TOKEN = os.environ.get("HF_TOKEN")', app_source)
768
- self.assertIn('token=ASSETS_HF_TOKEN or False', app_source)
769
- self.assertIn('HF_HUB_DISABLE_IMPLICIT_TOKEN', app_source)
770
-
771
-
772
- if __name__ == "__main__":
773
- unittest.main()
 
 
 
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()