"""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_stops_remaining_on_should_stop(self): """§19.24: should_stop 触发后不再提交剩余图片，返回「已停止」结果。""" from core.ai_outfit import generate_outfit_image d = self._make_dir_with_images("a") # 3 images, image_concurrency=1 out = self.tmp / "out" client = _RecordingClient(self._image_bytes()) # stop as soon as the first image has been generated def should_stop(): return client.calls >= 1 result = generate_outfit_image( self._dir_task(d), "话术", out, model_config={}, api_client=client, should_stop=should_stop) self.assertFalse(result.success) self.assertEqual(client.calls, 1) # only 1 generated; rest not submitted self.assertIn("已停止", result.error) self.assertIn("2 张未生成", result.error) self.assertEqual(len(result.output_paths), 1) 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()