cmbot/tests/test_ai_outfit.py

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