feat: add outfit generation core
This commit is contained in:
@@ -0,0 +1,123 @@
|
||||
"""Tests for single-row AI outfit generation core."""
|
||||
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())
|
||||
|
||||
self.assertEqual(prompt, "商品 纯棉短袖 / TY001")
|
||||
|
||||
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())
|
||||
self.assertEqual(client.prompt, "为 纯棉短袖 生成 TY001")
|
||||
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)
|
||||
|
||||
|
||||
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")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user