Excel C 列除单张图片外也可填一个目录(d:/images/a/):对目录内每张图各 生成一张穿搭图(N→N),全部存到 输出目录/<目录叶子名>/,文件名沿用源图名; Excel 仍整行一个状态:D=子目录、E=全成功才「完成」、F=失败张数/原因。 - core/ai_outfit.py: list_directory_images(顶层、扩展名过滤、排序、忽略子目录)、 make_outfit_subdir_path(不加 _n 后缀)、generate_outfit_image 目录分支 (新参 request_interval/image_log;逐张跳过已存在→幂等重试、本地节流、发日志、 聚合成单个 OutfitResult) - core/models.py: OutfitResult 加 output_paths(目录行各 jpg,供缩略图) - ai_outfit_panel.py: gen 闭包传 request_interval/image_log;缩略图逐张; 明细「结果」列显示「子目录(N 张)」;_basename 处理目录末尾分隔符 - tests/test_ai_outfit.py: 扇出/幂等跳过/空目录/缺目录/部分失败/命名过滤等用例 - docs/11 §4.1+§9.1、tasks.md §19.12 离屏冒烟:3 图目录 + mock API → output/<dir>/ 3 jpg、D=子目录、E=完成、缩略图 3 张; 全套 12 测试文件在 Python 3.7 全绿。 Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
298 lines
10 KiB
Python
298 lines
10 KiB
Python
"""Tests for single-row AI outfit generation core."""
|
|
import os
|
|
import shutil
|
|
import sys
|
|
import tempfile
|
|
import unittest
|
|
from io import BytesIO
|
|
from pathlib import Path
|
|
|
|
from PIL import Image
|
|
|
|
sys.path.insert(0, str(Path(__file__).parent.parent / "src"))
|
|
|
|
|
|
class TestAiOutfitCore(unittest.TestCase):
|
|
def setUp(self):
|
|
self.tmp = Path(tempfile.mkdtemp())
|
|
|
|
def tearDown(self):
|
|
shutil.rmtree(str(self.tmp), ignore_errors=True)
|
|
|
|
def _task(self, product_id="TY001"):
|
|
from core.models import OutfitTask
|
|
|
|
return OutfitTask(
|
|
row_index=2,
|
|
title="纯棉短袖",
|
|
product_id=product_id,
|
|
garment_path=str(self.tmp / "garment.png"),
|
|
)
|
|
|
|
def _image_bytes(self, size=(640, 480), color=(200, 120, 80)):
|
|
img = Image.new("RGB", size, color)
|
|
buffer = BytesIO()
|
|
img.save(buffer, format="PNG")
|
|
return buffer.getvalue()
|
|
|
|
def test_render_prompt_replaces_placeholders(self):
|
|
from core.ai_outfit import render_prompt
|
|
|
|
prompt = render_prompt("商品 {title} / {product_id}", self._task())
|
|
|
|
# No resolution -> no appended requirements tail.
|
|
self.assertEqual(prompt, "商品 纯棉短袖 / TY001")
|
|
|
|
def test_render_prompt_appends_requirements_with_resolution(self):
|
|
from core.ai_outfit import render_prompt
|
|
|
|
prompt = render_prompt("话术 {title}", self._task(), resolution="2K")
|
|
|
|
self.assertTrue(prompt.startswith("话术 纯棉短袖"))
|
|
self.assertIn("批量生成输出要求:", prompt)
|
|
self.assertIn("参考解析度:2K", prompt)
|
|
self.assertIn("不可跑版", prompt)
|
|
|
|
def test_build_output_requirements_empty_and_filled(self):
|
|
from core.ai_outfit import build_output_requirements
|
|
|
|
self.assertEqual(build_output_requirements(""), "")
|
|
self.assertEqual(build_output_requirements(None), "")
|
|
filled = build_output_requirements("4K")
|
|
self.assertIn("参考解析度:4K", filled)
|
|
self.assertIn("固定 1:1 正方形主图", filled)
|
|
|
|
def test_safe_product_filename_replaces_invalid_chars(self):
|
|
from core.ai_outfit import safe_product_filename
|
|
|
|
self.assertEqual(safe_product_filename('TY:00/1*?"'), "TY_00_1___")
|
|
|
|
def test_make_output_path_avoids_overwrite(self):
|
|
from core.ai_outfit import make_outfit_output_path
|
|
|
|
first = make_outfit_output_path(self.tmp, "TY001")
|
|
first.write_text("exists")
|
|
|
|
second = make_outfit_output_path(self.tmp, "TY001")
|
|
|
|
self.assertEqual(second.name, "TY001_1.jpg")
|
|
|
|
def test_save_jpg_under_limit_outputs_square_jpg(self):
|
|
from core.ai_outfit import save_jpg_under_limit
|
|
|
|
out = self.tmp / "out.jpg"
|
|
save_jpg_under_limit(self._image_bytes(size=(640, 480)), out, quality=85, max_bytes=200000)
|
|
|
|
self.assertTrue(out.exists())
|
|
self.assertLessEqual(out.stat().st_size, 200000)
|
|
with Image.open(str(out)) as img:
|
|
self.assertEqual(img.format, "JPEG")
|
|
self.assertEqual(img.size[0], img.size[1])
|
|
|
|
def test_generate_outfit_image_success(self):
|
|
from core.ai_outfit import generate_outfit_image
|
|
|
|
client = _FakeClient(self._image_bytes())
|
|
|
|
result = generate_outfit_image(
|
|
self._task(),
|
|
"为 {title} 生成 {product_id}",
|
|
self.tmp,
|
|
model_config={"url": "https://api", "model": "m", "api_key": "k"},
|
|
api_client=client,
|
|
)
|
|
|
|
self.assertTrue(result.success, result.error)
|
|
self.assertTrue(Path(result.output_path).exists())
|
|
# generate passes resolution -> rendered prompt + appended requirements.
|
|
self.assertTrue(client.prompt.startswith("为 纯棉短袖 生成 TY001"))
|
|
self.assertIn("批量生成输出要求:", client.prompt)
|
|
self.assertTrue(client.image_path.endswith("garment.png"))
|
|
|
|
def test_generate_outfit_image_failure(self):
|
|
from core.ai_outfit import generate_outfit_image
|
|
|
|
with self.assertLogs("core.ai_outfit", level="ERROR"):
|
|
result = generate_outfit_image(
|
|
self._task(),
|
|
"prompt",
|
|
self.tmp,
|
|
model_config={},
|
|
api_client=_FailingClient(),
|
|
)
|
|
|
|
self.assertFalse(result.success)
|
|
self.assertIn("boom", result.error)
|
|
|
|
# -- directory rows (docs/11 §4.1 / §9.1) ---------------------------
|
|
|
|
def _make_image_file(self, path, color=(10, 20, 30)):
|
|
Image.new("RGB", (64, 64), color).save(str(path), format="PNG")
|
|
|
|
def _dir_task(self, garment_path, product_id="DIRA"):
|
|
from core.models import OutfitTask
|
|
|
|
return OutfitTask(row_index=3, title="目录款", product_id=product_id,
|
|
garment_path=str(garment_path))
|
|
|
|
def _make_dir_with_images(self, name="a", files=("img1.png", "img2.png", "img3.png")):
|
|
d = self.tmp / name
|
|
d.mkdir()
|
|
for fname in files:
|
|
self._make_image_file(d / fname)
|
|
return d
|
|
|
|
def test_list_directory_images_filters_sorts_ignores_subdirs(self):
|
|
from core.ai_outfit import list_directory_images
|
|
|
|
d = self.tmp / "imgs"
|
|
d.mkdir()
|
|
self._make_image_file(d / "b.png")
|
|
self._make_image_file(d / "a.jpg")
|
|
(d / "note.txt").write_text("x", encoding="utf-8")
|
|
(d / "sub").mkdir()
|
|
self._make_image_file(d / "sub" / "c.png")
|
|
|
|
images = list_directory_images(d)
|
|
|
|
self.assertEqual([p.name for p in images], ["a.jpg", "b.png"])
|
|
|
|
def test_make_outfit_subdir_path_sanitizes_without_suffix(self):
|
|
from core.ai_outfit import make_outfit_subdir_path
|
|
|
|
p = make_outfit_subdir_path(self.tmp, "a:b", "img/1")
|
|
|
|
self.assertEqual(p.parent.name, "a_b")
|
|
self.assertEqual(p.name, "img_1.jpg")
|
|
|
|
def test_generate_directory_fans_out_to_named_subdir(self):
|
|
from core.ai_outfit import generate_outfit_image
|
|
|
|
d = self._make_dir_with_images("a")
|
|
out = self.tmp / "out"
|
|
client = _RecordingClient(self._image_bytes())
|
|
|
|
result = generate_outfit_image(
|
|
self._dir_task(d), "话术 {title}", out,
|
|
model_config={}, api_client=client,
|
|
)
|
|
|
|
self.assertTrue(result.success, result.error)
|
|
self.assertEqual(client.calls, 3)
|
|
self.assertEqual(Path(result.output_path), out / "a")
|
|
self.assertEqual(len(result.output_paths), 3)
|
|
names = sorted(p.name for p in (out / "a").iterdir())
|
|
self.assertEqual(names, ["img1.jpg", "img2.jpg", "img3.jpg"])
|
|
|
|
def test_generate_directory_skips_existing_outputs_on_retry(self):
|
|
from core.ai_outfit import generate_outfit_image
|
|
|
|
d = self._make_dir_with_images("a")
|
|
out = self.tmp / "out"
|
|
first = generate_outfit_image(
|
|
self._dir_task(d), "x {title}", out,
|
|
model_config={}, api_client=_RecordingClient(self._image_bytes()))
|
|
self.assertTrue(first.success)
|
|
|
|
# Re-run: every output already exists -> no API calls, still success.
|
|
again_client = _RecordingClient(self._image_bytes())
|
|
again = generate_outfit_image(
|
|
self._dir_task(d), "x {title}", out,
|
|
model_config={}, api_client=again_client)
|
|
|
|
self.assertTrue(again.success)
|
|
self.assertEqual(again_client.calls, 0)
|
|
self.assertEqual(len(again.output_paths), 3)
|
|
|
|
def test_generate_directory_empty_fails(self):
|
|
from core.ai_outfit import generate_outfit_image
|
|
|
|
d = self.tmp / "empty"
|
|
d.mkdir()
|
|
|
|
result = generate_outfit_image(
|
|
self._dir_task(d), "x", self.tmp / "out",
|
|
model_config={}, api_client=_RecordingClient(self._image_bytes()))
|
|
|
|
self.assertFalse(result.success)
|
|
self.assertIn("没有图片", result.error)
|
|
|
|
def test_generate_directory_missing_fails(self):
|
|
from core.ai_outfit import generate_outfit_image
|
|
|
|
# Trailing separator marks it as a directory even though it doesn't exist.
|
|
missing = str(self.tmp / "nope") + os.sep
|
|
result = generate_outfit_image(
|
|
self._dir_task(missing), "x", self.tmp / "out",
|
|
model_config={}, api_client=_RecordingClient(self._image_bytes()))
|
|
|
|
self.assertFalse(result.success)
|
|
self.assertIn("目录不存在", result.error)
|
|
|
|
def test_generate_directory_partial_failure_aggregates(self):
|
|
from core.ai_outfit import generate_outfit_image
|
|
|
|
d = self._make_dir_with_images("a")
|
|
out = self.tmp / "out"
|
|
client = _FailOnClient(self._image_bytes(), fail_name="img2.png")
|
|
|
|
result = generate_outfit_image(
|
|
self._dir_task(d), "x", out, model_config={}, api_client=client)
|
|
|
|
self.assertFalse(result.success)
|
|
self.assertIn("3 张中 1 张失败", result.error)
|
|
self.assertIn("img2.png", result.error)
|
|
# The two that succeeded are still written (and listed for thumbnails).
|
|
self.assertEqual(len(result.output_paths), 2)
|
|
self.assertFalse((out / "a" / "img2.jpg").exists())
|
|
|
|
|
|
class _FakeClient:
|
|
def __init__(self, image_bytes):
|
|
self._image_bytes = image_bytes
|
|
self.prompt = None
|
|
self.image_path = None
|
|
|
|
def generate(self, prompt, image_path, resolution="1K"):
|
|
self.prompt = prompt
|
|
self.image_path = str(image_path)
|
|
return self._image_bytes
|
|
|
|
|
|
class _FailingClient:
|
|
def generate(self, prompt, image_path, resolution="1K"):
|
|
raise RuntimeError("boom")
|
|
|
|
|
|
class _RecordingClient:
|
|
"""Records every generate() call (count + image paths) for directory tests."""
|
|
|
|
def __init__(self, image_bytes):
|
|
self._image_bytes = image_bytes
|
|
self.calls = 0
|
|
self.image_paths = []
|
|
|
|
def generate(self, prompt, image_path, resolution="1K"):
|
|
self.calls += 1
|
|
self.image_paths.append(str(image_path))
|
|
return self._image_bytes
|
|
|
|
|
|
class _FailOnClient:
|
|
"""Fails only for the source image whose filename ends with *fail_name*."""
|
|
|
|
def __init__(self, image_bytes, fail_name):
|
|
self._image_bytes = image_bytes
|
|
self._fail_name = fail_name
|
|
self.calls = 0
|
|
|
|
def generate(self, prompt, image_path, resolution="1K"):
|
|
self.calls += 1
|
|
if str(image_path).endswith(self._fail_name):
|
|
raise RuntimeError("bad image")
|
|
return self._image_bytes
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|