Agent/examples/auto_put_ad_mini/test_ai_generated_materials.py

import unittest
import os
import sys
from pathlib import Path

from dotenv import load_dotenv

_HERE = Path(__file__).parent
sys.path.insert(0, str(_HERE))
sys.path.insert(0, str(_HERE.parent.parent))

from tools.ai_generated_material import (
    _public_oss_url,
    build_generation_prompts,
)
from tools.account_material_strategy import normalize_material_source, parse_bool_flag
from tools.video_feature_query import read_cached_video_element_features


class AiGeneratedMaterialsTest(unittest.TestCase):
    def test_normalize_material_source_defaults_to_history(self):
        self.assertEqual(normalize_material_source(""), "history")
        self.assertEqual(normalize_material_source("历史素材"), "history")
        self.assertEqual(normalize_material_source("AI生成素材"), "ai_generated")
        self.assertEqual(normalize_material_source("外部素材"), "external_recall")

    def test_parse_bool_flag_accepts_feishu_yes_no(self):
        self.assertTrue(parse_bool_flag("是"))
        self.assertFalse(parse_bool_flag("否"))
        self.assertFalse(parse_bool_flag(""))

    def test_build_generation_prompts_uses_cached_video_topic(self):
        load_dotenv(_HERE / ".env")
        video_id = int(os.getenv("AI_MATERIAL_TEST_VIDEO_ID", "69131739"))
        try:
            features = read_cached_video_element_features([video_id]).get(video_id) or []
        except Exception as e:
            self.skipTest(f"数据库不可用,跳过真实视频特征测试:{e}")
        topic = next((f for f in features if f.element_dimension == "解构选题"), None)
        if topic is None:
            self.skipTest(f"video_id={video_id} 在本地缓存表中没有解构选题")

        prompts = build_generation_prompts(
            video_id=video_id,
            title="不应作为兜底标题",
            category="不应作为兜底品类",
            features=features,
        )

        self.assertEqual([p.prompt_type for p in prompts], ["topic"])
        self.assertIn(topic.standard_element, prompts[0].prompt_text)
        self.assertEqual(prompts[0].feature_hits[0]["original_standard_element"], topic.standard_element)
        self.assertEqual(prompts[0].feature_hits[0]["video_description"], topic.standard_element)
        self.assertIn("腾讯广告信息流 / 公众号投放", prompts[0].prompt_text)
        self.assertNotIn("不应作为兜底标题", prompts[0].prompt_text)
        self.assertNotIn("不应作为兜底品类", prompts[0].prompt_text)
        self.assertIn("16:9", prompts[0].prompt_text)

    def test_public_oss_url_keeps_directory_slashes(self):
        self.assertEqual(
            _public_oss_url(
                "https://rescdn.yishihui.com/",
                "auto_put_tencent/image/a b/test.jpg",
            ),
            "https://rescdn.yishihui.com/auto_put_tencent/image/a%20b/test.jpg",
        )

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