cmshoppe/tests/test_prompts.py

import os
import sys
import unittest

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

from _helpers import TempDirMixin

from app import prompts


class PromptTests(TempDirMixin, unittest.TestCase):
    def test_title_prompt_load_missing_and_save(self):
        with self.make_temp_dir() as temp_dir:
            path = os.path.join(temp_dir, "title_prompt.txt")

            self.assertEqual("", prompts.load_title_prompt(path))
            prompts.save_title_prompt("标题规则", path)

            self.assertEqual("标题规则", prompts.load_title_prompt(path))

        self.assert_removed(temp_dir)

    def test_cover_template_crud_and_validation(self):
        with self.make_temp_dir() as temp_dir:
            directory = os.path.join(temp_dir, "prompts", "cover")

            self.assertEqual([], prompts.list_cover_templates(directory))
            prompts.save_cover_template("韩版女装", "封面规则", directory)
            prompts.save_cover_template("基础", "基础规则", directory)

            self.assertEqual(["基础", "韩版女装"], prompts.list_cover_templates(directory))
            self.assertEqual("封面规则", prompts.load_cover_template("韩版女装", directory))

            prompts.rename_cover_template("基础", "基础2", directory)
            self.assertEqual(["基础2", "韩版女装"], prompts.list_cover_templates(directory))

            with self.assertRaises(prompts.PromptError):
                prompts.rename_cover_template("基础2", "韩版女装", directory)
            with self.assertRaises(prompts.PromptError):
                prompts.save_cover_template("../bad", "x", directory)

            prompts.delete_cover_template("基础2", directory)
            self.assertEqual(["韩版女装"], prompts.list_cover_templates(directory))

        self.assert_removed(temp_dir)

    def test_ensure_default_prompts_seeds_empty_user_prompt_files(self):
        with self.make_temp_dir() as temp_dir:
            title_path = os.path.join(temp_dir, "title_prompt.txt")
            cover_dir = os.path.join(temp_dir, "prompts", "cover")

            prompts.ensure_default_prompts(title_path, cover_dir)

            self.assertIn("蝦皮台灣站", prompts.load_title_prompt(title_path))
            self.assertEqual(["papa1"], prompts.list_cover_templates(cover_dir))
            self.assertIn(
                "商品标题:{新标题}",
                prompts.load_cover_template("papa1", cover_dir),
            )

        self.assert_removed(temp_dir)

    def test_ensure_default_prompts_does_not_overwrite_user_prompts(self):
        with self.make_temp_dir() as temp_dir:
            title_path = os.path.join(temp_dir, "title_prompt.txt")
            cover_dir = os.path.join(temp_dir, "prompts", "cover")
            prompts.save_title_prompt("用户标题提示词", title_path)
            prompts.save_cover_template("用户模板", "用户封面提示词", cover_dir)

            prompts.ensure_default_prompts(title_path, cover_dir)

            self.assertEqual("用户标题提示词", prompts.load_title_prompt(title_path))
            self.assertEqual(["用户模板"], prompts.list_cover_templates(cover_dir))
            self.assertEqual(
                "用户封面提示词",
                prompts.load_cover_template("用户模板", cover_dir),
            )

        self.assert_removed(temp_dir)

    def test_render_prompt_replaces_known_variables(self):
        task = {
            "old_title": "舊T恤",
            "new_title": "新T恤",
            "item_id": "51100639510",
            "account_name": "主店",
        }

        rendered = prompts.render_prompt(
            "用{旧标题}生成{新标题}，商品{商品id}，店铺{店铺}，未知{不存在}",
            task,
        )

        self.assertEqual(
            "用舊T恤生成新T恤，商品51100639510，店铺主店，未知{不存在}",
            rendered,
        )


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