109 lines
3.5 KiB
Python
109 lines
3.5 KiB
Python
 
|
"""Tests for AI title generation core (docs/11 §17)."""
|
||
|
|
import shutil
|
||
|
|
import sys
|
||
|
|
import tempfile
|
||
|
|
import unittest
|
||
|
|
from pathlib import Path
|
||
|
|
|
||
|
|
from PIL import Image
|
||
|
|
|
||
|
|
sys.path.insert(0, str(Path(__file__).parent.parent / "src"))
|
||
|
|
|
||
|
|
|
||
|
|
class TestAiTitleCore(unittest.TestCase):
|
||
|
|
def setUp(self):
|
||
|
|
self.tmp = Path(tempfile.mkdtemp())
|
||
|
|
|
||
|
|
def tearDown(self):
|
||
|
|
shutil.rmtree(str(self.tmp), ignore_errors=True)
|
||
|
|
|
||
|
|
def _make_image(self, path):
|
||
|
|
Image.new("RGB", (32, 32), (120, 80, 40)).save(str(path), format="PNG")
|
||
|
|
|
||
|
|
def _task(self, garment_path, title="占位标题", product_id=""):
|
||
|
|
from core.models import OutfitTask
|
||
|
|
|
||
|
|
return OutfitTask(row_index=2, title=title, product_id=product_id,
|
||
|
|
garment_path=str(garment_path))
|
||
|
|
|
||
|
|
def test_render_title_prompt_replaces_placeholders(self):
|
||
|
|
from core.ai_title import render_title_prompt
|
||
|
|
|
||
|
|
task = self._task("g.png", title="旧标题", product_id="TY001")
|
||
|
|
out = render_title_prompt("参考 {title} / {product_id}", task)
|
||
|
|
|
||
|
|
self.assertEqual(out, "参考 旧标题 / TY001")
|
||
|
|
|
||
|
|
def test_generate_title_single_file_success(self):
|
||
|
|
from core.ai_title import generate_title
|
||
|
|
|
||
|
|
garment = self.tmp / "shirt.png"
|
||
|
|
self._make_image(garment)
|
||
|
|
client = _RecordingTextClient("纯棉短袖T恤")
|
||
|
|
|
||
|
|
result = generate_title(self._task(garment), "起个标题 {title}",
|
||
|
|
model_config={}, api_client=client)
|
||
|
|
|
||
|
|
self.assertTrue(result.success, result.error)
|
||
|
|
self.assertEqual(result.generated_title, "纯棉短袖T恤")
|
||
|
|
self.assertEqual(client.image_paths, [str(garment)])
|
||
|
|
|
||
|
|
def test_generate_title_directory_uses_first_image(self):
|
||
|
|
from core.ai_title import generate_title
|
||
|
|
|
||
|
|
d = self.tmp / "FG201"
|
||
|
|
d.mkdir()
|
||
|
|
for name in ("b.png", "a.png", "c.png"):
|
||
|
|
self._make_image(d / name)
|
||
|
|
client = _RecordingTextClient("印花连衣裙")
|
||
|
|
|
||
|
|
result = generate_title(self._task(str(d) + "/"), "话术",
|
||
|
|
model_config={}, api_client=client)
|
||
|
|
|
||
|
|
self.assertTrue(result.success, result.error)
|
||
|
|
# Sorted: a.png is the first reference image.
|
||
|
|
self.assertEqual(client.image_paths, [str(d / "a.png")])
|
||
|
|
|
||
|
|
def test_generate_title_empty_directory_fails(self):
|
||
|
|
from core.ai_title import generate_title
|
||
|
|
|
||
|
|
d = self.tmp / "empty"
|
||
|
|
d.mkdir()
|
||
|
|
|
||
|
|
result = generate_title(self._task(str(d) + "/"), "话术",
|
||
|
|
model_config={}, api_client=_RecordingTextClient("x"))
|
||
|
|
|
||
|
|
self.assertFalse(result.success)
|
||
|
|
self.assertIn("没有图片", result.error)
|
||
|
|
|
||
|
|
def test_generate_title_client_error_aggregated(self):
|
||
|
|
from core.ai_title import generate_title
|
||
|
|
|
||
|
|
garment = self.tmp / "shirt.png"
|
||
|
|
self._make_image(garment)
|
||
|
|
|
||
|
|
result = generate_title(self._task(garment), "话术",
|
||
|
|
model_config={}, api_client=_FailingTextClient())
|
||
|
|
|
||
|
|
self.assertFalse(result.success)
|
||
|
|
self.assertIn("boom", result.error)
|
||
|
|
|
||
|
|
|
||
|
|
class _RecordingTextClient:
|
||
|
|
def __init__(self, title):
|
||
|
|
self._title = title
|
||
|
|
self.image_paths = []
|
||
|
|
|
||
|
|
def generate_text(self, prompt, image_path=None, resolution="1K"):
|
||
|
|
self.image_paths.append(str(image_path))
|
||
|
|
return self._title
|
||
|
|
|
||
|
|
|
||
|
|
class _FailingTextClient:
|
||
|
|
def generate_text(self, prompt, image_path=None, resolution="1K"):
|
||
|
|
raise RuntimeError("boom")
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
unittest.main()
|