TravelContentCreator/examples/test_workflow.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
测试脚本：验证旅游内容创作工具的完整工作流程
"""
import os
import sys
import json
import time

# 添加项目根目录到Python路径
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

# 导入所需模块
from main import load_config, generate_topics_step, generate_content_and_posters_step
from utils.tweet_generator import run_topic_generation_pipeline
from core.topic_parser import TopicParser

def test_full_workflow():
    """测试完整的工作流程，从选题生成到海报制作"""
    print("测试完整工作流程...")
    
    try:
        # 1. 加载配置
        print("步骤 1: 加载配置...")
        config = load_config()
        print(f"配置加载成功: {len(config)} 个配置项")
        
        # 2. 执行选题生成
        print("\n步骤 2: 生成选题...")
        run_id, tweet_topic_record = run_topic_generation_pipeline(config)
        
        if not run_id or not tweet_topic_record:
            print("选题生成失败，测试终止。")
            return False
        
        print(f"选题生成成功，Run ID: {run_id}")
        print(f"生成了 {len(tweet_topic_record.topics_list)} 个选题")
        
        # 3. 执行内容和海报生成
        print("\n步骤 3: 生成内容和海报...")
        generate_content_and_posters_step(config, run_id, tweet_topic_record)
        
        # 4. 验证输出
        output_dir = os.path.join(config["output_dir"], run_id)
        if os.path.exists(output_dir):
            items = os.listdir(output_dir)
            print(f"\n输出目录 {output_dir} 中有 {len(items)} 个项目")
            
            # 检查是否有生成的海报
            poster_count = 0
            for item in items:
                if item.startswith(("1_", "2_", "3_", "4_", "5_")):  # 假设选题编号从1开始
                    poster_dir = os.path.join(output_dir, item, "poster")
                    if os.path.exists(poster_dir):
                        poster_files = [f for f in os.listdir(poster_dir) if f.endswith(".jpg")]
                        poster_count += len(poster_files)
            
            print(f"共生成 {poster_count} 个海报")
            return poster_count > 0
        else:
            print(f"输出目录 {output_dir} 不存在")
            return False
    
    except Exception as e:
        print(f"测试过程中出错: {e}")
        import traceback
        traceback.print_exc()
        return False

def test_steps_separately():
    """分步测试工作流程"""
    print("分步测试工作流程...")
    
    try:
        # 1. 加载配置
        print("步骤 1: 加载配置...")
        config = load_config()
        
        # 将variants设置为较小的值以加快测试
        test_config = config.copy()
        test_config["variants"] = 1  # 每个选题只生成1个变体
        test_config["num"] = 2  # 只生成2个选题
        
        print(f"测试配置已准备: variants={test_config['variants']}, num={test_config['num']}")
        
        # 2. 仅执行选题生成
        print("\n步骤 2: 仅测试选题生成...")
        run_id, tweet_topic_record = run_topic_generation_pipeline(test_config)
        
        if not run_id or not tweet_topic_record:
            print("选题生成失败，测试终止。")
            return False
        
        # 保存生成的选题数据用于后续测试
        topics_file = os.path.join(test_config["output_dir"], run_id, "tweet_topic.json")
        print(f"选题结果已保存到: {topics_file}")
        
        # 3. 模拟间隔后继续处理
        print("\n步骤 3: 模拟间隔后处理选题结果...")
        print("（在真实场景中，可能是在不同的时间或机器上进行后续处理）")
        time.sleep(2)  # 模拟短暂延迟
        
        # 4. 从保存的选题文件加载数据，并手动执行内容和海报生成
        print("\n步骤 4: 加载保存的选题，执行内容和海报生成...")
        
        # 这部分通常是由main函数中的流程自动处理的
        # 这里为了演示分段流程，模拟手动加载数据并处理
        
        if os.path.exists(topics_file):
            with open(topics_file, 'r', encoding='utf-8') as f:
                topics_data = json.load(f)
            
            loaded_topic_record = TopicParser.from_json(topics_data)
            print(f"从 {topics_file} 加载了 {len(loaded_topic_record.topics_list)} 个选题")
            
            # 执行内容和海报生成步骤
            generate_content_and_posters_step(test_config, run_id, loaded_topic_record)
            
            return True
        else:
            print(f"选题文件 {topics_file} 不存在")
            return False
    
    except Exception as e:
        print(f"测试过程中出错: {e}")
        import traceback
        traceback.print_exc()
        return False

if __name__ == "__main__":
    print("==== 旅游内容创作工具工作流程测试 ====\n")
    
    # 判断测试模式
    import argparse
    parser = argparse.ArgumentParser(description='测试旅游内容创作工具的工作流程')
    parser.add_argument('--mode', type=str, choices=['full', 'steps', 'both'], 
                      default='both', help='测试模式: full=完整流程, steps=分步测试, both=两种模式都测试')
    args = parser.parse_args()
    
    success = True
    
    if args.mode in ['full', 'both']:
        print("\n==== 测试完整工作流程 ====")
        success = test_full_workflow() and success
    
    if args.mode in ['steps', 'both']:
        print("\n==== 测试分步工作流程 ====")
        success = test_steps_separately() and success
    
    if success:
        print("\n==== 所有测试通过! ====")
        sys.exit(0)
    else:
        print("\n==== 测试失败! ====")
        sys.exit(1)