cmshoppe/tests/test_product_suite.py

import os
import sys
import unittest
from types import SimpleNamespace

sys.path.insert(0, os.path.dirname(__file__))

from app import product_suite, prompts


class ProductSuiteTests(unittest.TestCase):
    def test_defaults_and_per_image_total(self):
        settings = product_suite.default_suite_settings()

        self.assertEqual(
            ["白底图", "场景图", "模特场景图", "细节说明图", "卖点图"],
            product_suite.category_order(settings),
        )
        self.assertEqual(5, product_suite.suite_total_count(settings, 3))

        settings["per_image_primary"] = True

        self.assertEqual(13, product_suite.suite_total_count(settings, 3))

    def test_custom_category_validation_and_normalization(self):
        self.assertEqual("分类名称不能为空", product_suite.suite_name_error(""))
        self.assertEqual("分类名称不能包含空格", product_suite.suite_name_error("使用 场景"))
        self.assertEqual("分类名称不能超过10个字", product_suite.suite_name_error("一二三四五六七八九十甲"))
        self.assertEqual("分类名称已存在", product_suite.suite_name_error("场景图", {"场景图"}))

        settings = product_suite.normalize_suite_settings(
            {
                "ratio": "3:4",
                "categories": {"白底图": 2, "场景图": 0, "卖点图": 1, "尺寸图": 2},
                "custom_category_order": ["尺寸图"],
            }
        )

        self.assertEqual("3:4", settings["ratio"])
        self.assertEqual(
            ["白底图", "场景图", "模特场景图", "细节说明图", "卖点图", "尺寸图"],
            product_suite.category_order(settings),
        )
        self.assertEqual(5, product_suite.suite_total_count(settings, 1))

    def test_recent_settings_only_keep_allowed_dropdown_values(self):
        self.assertEqual(
            {
                "platform": "Shopee",
                "country": "中国台湾",
                "language": "繁体中文",
                "ratio": "1:1",
            },
            product_suite.last_suite_settings(
                {
                    "platform": "未知平台",
                    "country": "未知站点",
                    "language": "未知语言",
                    "ratio": "2:3",
                    "per_image_primary": True,
                    "categories": {"白底图": 9},
                }
            ),
        )
        selected = product_suite.last_suite_settings(
            {
                "platform": "Amazon",
                "country": "新加坡",
                "language": "英文",
                "ratio": "4:3",
            }
        )
        self.assertEqual("Amazon", selected["platform"])
        self.assertEqual("新加坡", selected["country"])
        self.assertEqual("英文", selected["language"])
        self.assertEqual("4:3", selected["ratio"])

    def test_job_specs_include_selected_context_and_source_assignment(self):
        settings = product_suite.default_suite_settings()
        settings.update(
            {
                "platform": "Shopee",
                "country": "中国台湾",
                "language": "繁体中文",
                "ratio": "4:3",
                "per_image_primary": True,
                "categories": {"白底图": 1, "场景图": 1, "卖点图": 0},
            }
        )
        assets = [SimpleNamespace(id=11), SimpleNamespace(id=12)]

        specs = product_suite.build_job_specs(
            assets,
            "40小时续航，适合通勤",
            settings,
            "51100639510",
            template_text=prompts.load_default_product_suite_prompt(),
        )

        self.assertEqual(3, len(specs))
        self.assertEqual([11, 11, 12], [spec["source_asset_id"] for spec in specs])
        self.assertEqual([[], [], []], [spec["reference_asset_ids"] for spec in specs])
        self.assertEqual(["白底图", "场景图", "场景图"], [spec["job_type"] for spec in specs])
        for spec in specs:
            self.assertIn("平台：Shopee", spec["prompt"])
            self.assertIn("国家地区：中国台湾", spec["prompt"])
            self.assertIn("输出语言：繁体中文", spec["prompt"])
            self.assertIn("本次生成比例：4:3", spec["prompt"])
            self.assertNotIn("商品ID：51100639510", spec["prompt"])
            self.assertIn("40小时续航", spec["prompt"])
            self.assertIn("禁止海报长图、详情页长图和多宫格拼接版面", spec["prompt"])
            self.assertIn("禁止在画面中出现任何国旗", spec["prompt"])
            self.assertIn("参考图规则:当前上传图片是本任务唯一主参考图", spec["prompt"])
            self.assertNotIn("参考图规则：", spec["prompt"])
            self.assertNotIn("序号%d" % spec["source_index"], spec["prompt"])
            self.assertTrue(
                spec["prompt"].splitlines()[-1].startswith("本次生成比例：4:3")
            )

        self.assertIn(
            "生成目标：白底图,生成 Shopee 台灣商品白底主圖，純白背景，商品清晰置中，不添加多餘文字。",
            specs[0]["prompt"],
        )
        self.assertIn(
            "生成目标：场景图,生成生活化使用場景圖，展示商品在真實情境中的用途與氛圍，畫面自然可信。",
            specs[1]["prompt"],
        )
        white_prompt = specs[0]["prompt"]
        reference_rule = (
            "当前上传图片是本任务唯一主参考图；保持商品主体、款式、颜色和关键细节准确；"
            "不编造用户与参考图均未提供的信息。"
        )
        self.assertLess(white_prompt.index("重要尺寸要求"), white_prompt.index("重要禁用内容"))
        self.assertLess(white_prompt.index("重要禁用内容"), white_prompt.index("价格信息规则"))
        self.assertLess(white_prompt.index("价格信息规则"), white_prompt.index("尺码信息规则"))
        self.assertLess(white_prompt.index("尺码信息规则"), white_prompt.index(reference_rule))
        self.assertLess(white_prompt.index(reference_rule), white_prompt.index("商品卖点与要求"))

    def test_job_specs_freeze_first_image_references_when_not_per_image_primary(self):
        settings = product_suite.default_suite_settings()
        settings.update(
            {
                "per_image_primary": False,
                "categories": {"白底图": 1, "场景图": 1, "卖点图": 0},
            }
        )
        specs = product_suite.build_job_specs(
            [SimpleNamespace(id=11), SimpleNamespace(id=12), SimpleNamespace(id=13)],
            "卖点",
            settings,
            "51100639510",
            template_text=prompts.load_default_product_suite_prompt(),
        )

        self.assertEqual(2, len(specs))
        self.assertEqual([11, 11], [spec["source_asset_id"] for spec in specs])
        self.assertEqual([[12, 13], [12, 13]], [spec["reference_asset_ids"] for spec in specs])

    def test_job_specs_use_fixed_reference_rule_for_one_three_and_eight_images(self):
        settings = product_suite.default_suite_settings()
        settings.update(
            {
                "per_image_primary": False,
                "categories": {"白底图": 1, "场景图": 0, "卖点图": 0},
            }
        )
        fixed_rule = "参考图规则:使用第一張上傳圖作為主商品圖，其餘圖片只作為參考。"

        for image_count in (1, 3, 8):
            specs = product_suite.build_job_specs(
                [SimpleNamespace(id=index) for index in range(1, image_count + 1)],
                "卖点",
                settings,
                "51100639510",
                template_text=prompts.load_default_product_suite_prompt(),
            )
            self.assertEqual(1, len(specs))
            self.assertIn(fixed_rule, specs[0]["prompt"])
            self.assertNotIn("第2至", specs[0]["prompt"])

    def test_reference_rule_follows_per_image_primary_setting(self):
        template = prompts.load_default_product_suite_prompt()
        settings = product_suite.default_suite_settings()

        fixed_rule = "参考图规则:使用第一張上傳圖作為主商品圖，其餘圖片只作為參考。"
        for reference_count in (0, 2, 7):
            not_per_image_primary = product_suite.build_suite_prompt(
                "卖点",
                settings,
                "场景图",
                "51100639510",
                reference_asset_count=reference_count,
                template_text=template,
            )
            self.assertIn(fixed_rule, not_per_image_primary)
            self.assertNotIn("第2至", not_per_image_primary)

        settings["per_image_primary"] = True
        per_image_primary = product_suite.build_suite_prompt(
            "卖点",
            settings,
            "场景图",
            "51100639510",
            source_index=2,
            template_text=template,
        )
        self.assertIn(
            "参考图规则:当前上传图片是本任务唯一主参考图；保持商品主体、款式、颜色和关键细节准确；不编造用户与参考图均未提供的信息。",
            per_image_primary,
        )
        self.assertNotIn("第2至", per_image_primary)
        self.assertNotIn("序号2", per_image_primary)

        self.assertEqual(7, product_suite.reference_asset_count({}, 9))

    def test_product_suite_template_validation_and_custom_category_rendering(self):
        default_text = prompts.load_default_product_suite_prompt()
        self.assertEqual([], product_suite.product_suite_prompt_errors(default_text))

        unknown = default_text + "\n{未知变量}"
        self.assertIn("未知变量", "；".join(product_suite.product_suite_prompt_errors(unknown)))

        missing = default_text.replace("{图片比例}", "")
        self.assertIn(
            "缺少必需变量",
            "；".join(product_suite.product_suite_prompt_errors(missing)),
        )

        inline_rule = default_text.replace(
            "{价格信息规则}",
            "价格：{价格信息规则}",
        )
        self.assertIn(
            "只读规则变量必须独占一行",
            "；".join(product_suite.product_suite_prompt_errors(inline_rule)),
        )

        literal_brace = default_text + "\n普通内容{"
        self.assertIn(
            "不支持的字面花括号",
            "；".join(product_suite.product_suite_prompt_errors(literal_brace)),
        )

        settings = product_suite.default_suite_settings()
        context = product_suite.product_suite_prompt_context(
            "突出轻量材质",
            settings,
            "尺寸图",
            "draft_hidden",
        )
        rendered = product_suite.render_product_suite_prompt(default_text, context)
        self.assertIn("生成目标：尺寸图,生成自定义分类图片：尺寸图。", rendered)
        self.assertNotIn("draft_", rendered)
        self.assertNotIn("套图名称：", rendered)

        legacy_template = default_text.replace(
            "生成目标：{套图名称},{生成目标}",
            "套图名称：{套图名称}{补充描述}",
        )
        self.assertEqual([], product_suite.product_suite_prompt_errors(legacy_template))
        legacy_rendered = product_suite.render_product_suite_prompt(legacy_template, context)
        self.assertIn("套图名称：尺寸图，生成自定义分类图片：尺寸图。", legacy_rendered)

    def test_category_descriptions_are_fixed_and_custom_name_is_explicit(self):
        expected = {
            "白底图": "生成 Shopee 台灣商品白底主圖，純白背景，商品清晰置中，不添加多餘文字。",
            "场景图": "生成生活化使用場景圖，展示商品在真實情境中的用途與氛圍，畫面自然可信。",
            "模特场景图": "生成模特或手持使用情境圖，畫面自然可信，商品為主角。",
            "细节说明图": "生成商品細節特寫說明圖，突出材質、做工、接口、紋理或關鍵細節。",
            "卖点图": "生成賣點詳解圖，使用繁體中文短文案呈現核心優勢，版面乾淨。",
        }
        for category, description in expected.items():
            self.assertEqual(description, product_suite.category_description(category))

        self.assertEqual(
            "生成自定义分类图片：结构展示。",
            product_suite.category_description("结构展示"),
        )

    def test_product_suite_optional_item_and_reference_variables(self):
        template = prompts.load_default_product_suite_prompt() + (
            "\n商品ID：{商品ID}\n参考图序号：{主参考图序号}"
        )
        context = product_suite.product_suite_prompt_context(
            "卖点",
            product_suite.default_suite_settings(),
            "卖点图",
            "",
            source_index=2,
        )

        rendered = product_suite.render_product_suite_prompt(template, context)

        self.assertIn("商品ID：未绑定商品", rendered)
        self.assertIn("参考图序号：2", rendered)


if __name__ == "__main__":
    unittest.main()