"""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) 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 / "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()