cmbot/tests/test_ai_outfit.py

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

        # 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)


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