353 lines
12 KiB
Python
353 lines
12 KiB
Python
"""Tests for single-row AI outfit generation core."""
|
|
import os
|
|
import shutil
|
|
import sys
|
|
import tempfile
|
|
import threading
|
|
import time
|
|
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)
|
|
|
|
def test_generate_single_file_empty_product_id_uses_source_name(self):
|
|
from core.ai_outfit import generate_outfit_image
|
|
from core.models import OutfitTask
|
|
|
|
garment = self.tmp / "shirt.png"
|
|
self._make_image_file(garment)
|
|
task = OutfitTask(row_index=2, title="款", product_id="",
|
|
garment_path=str(garment))
|
|
|
|
result = generate_outfit_image(
|
|
task, "穿 {title}", self.tmp / "o",
|
|
model_config={}, api_client=_RecordingClient(self._image_bytes()))
|
|
|
|
self.assertTrue(result.success, result.error)
|
|
self.assertEqual(Path(result.output_path).name, "shirt.jpg")
|
|
|
|
# -- 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 / "穿搭图片"
|
|
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())
|
|
|
|
def test_generate_directory_uses_image_concurrency(self):
|
|
from core.ai_outfit import generate_outfit_image
|
|
|
|
d = self._make_dir_with_images("a", files=("img1.png", "img2.png", "img3.png", "img4.png"))
|
|
out = self.tmp / "out"
|
|
client = _ConcurrentRecordingClient(self._image_bytes())
|
|
|
|
result = generate_outfit_image(
|
|
self._dir_task(d), "x", out, model_config={}, api_client=client,
|
|
image_concurrency=2)
|
|
|
|
self.assertTrue(result.success, result.error)
|
|
self.assertEqual(client.calls, 4)
|
|
self.assertGreaterEqual(client.max_active, 2)
|
|
self.assertEqual(len(result.output_paths), 4)
|
|
|
|
|
|
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
|
|
|
|
|
|
class _ConcurrentRecordingClient:
|
|
"""Thread-safe fake client that records concurrent generate() calls."""
|
|
|
|
def __init__(self, image_bytes):
|
|
self._image_bytes = image_bytes
|
|
self.calls = 0
|
|
self.active = 0
|
|
self.max_active = 0
|
|
self._lock = threading.Lock()
|
|
|
|
def generate(self, prompt, image_path, resolution="1K"):
|
|
with self._lock:
|
|
self.calls += 1
|
|
self.active += 1
|
|
self.max_active = max(self.max_active, self.active)
|
|
time.sleep(0.05)
|
|
with self._lock:
|
|
self.active -= 1
|
|
return self._image_bytes
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|