test_full_workflow.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
完整工作流测试脚本
测试从文本到视频的完整转换流程
"""

import os
import sys
import time
import logging
import unittest
from datetime import datetime

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)

# 添加项目根目录到路径
script_dir = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, script_dir)

from src.models import Session, Task, init_db
from src.processor import TaskProcessor

class FullWorkflowTest(unittest.TestCase):
    """测试完整工作流程"""
    
    def setUp(self):
        """测试前准备"""
        # 初始化数据库
        init_db()
        
        # 创建输出目录
        self.output_dir = os.path.join(script_dir, 'output', f'test_{datetime.now().strftime("%Y%m%d_%H%M%S")}')
        os.makedirs(self.output_dir, exist_ok=True)
        
        # 创建任务处理器
        self.task_processor = TaskProcessor()
        
        # 创建数据库会话
        self.session = Session()
    
    def tearDown(self):
        """测试后清理"""
        # 关闭数据库会话
        self.session.close()
    
    def test_create_task(self):
        """测试创建任务"""
        # 创建测试任务
        task = Task(
            name="测试任务",
            topic="科技发展",
            requirements="未来科技发展趋势，包括AI、机器人等领域",
            prompt="",  # 由任务处理器生成
            image_width=1024,
            image_height=1024,
            image_count=1,
            video_size="1280x720",  # 使用字符串格式的视频尺寸
            fps=30,
            with_audio=True,
            background_music=None,
            output_folder=self.output_dir
        )
        
        # 保存任务
        self.session.add(task)
        self.session.commit()
        
        # 验证任务已创建
        self.assertIsNotNone(task.id, "任务创建失败")
        logging.info(f"创建任务成功，ID: {task.id}")
        
        return task
    
    def test_full_workflow(self):
        """测试完整工作流程"""
        # 创建任务
        task = self.test_create_task()
        
        # 启动任务处理
        self.task_processor.start_task(task.id)
        
        # 监控任务状态
        timeout = 600  # 10分钟超时
        start_time = time.time()
        completed = False
        
        logging.info("开始监控任务状态...")
        
        while time.time() - start_time < timeout:
            # 刷新任务状态
            self.session.refresh(task)
            
            # 记录当前状态
            status_text = {
                0: "待处理",
                1: "生成提示词中",
                2: "生成图片中",
                3: "生成视频中", 
                4: "剪辑中",
                5: "已完成",
                -1: "失败"
            }.get(task.status, "未知")
            
            logging.info(f"任务 {task.id} 状态: {status_text}, 进度: {task.progress:.1%}, 阶段: {task.current_stage}")
            
            # 检查是否完成或失败
            if task.status == 5:  # 已完成
                completed = True
                break
            elif task.status == -1:  # 失败
                self.fail(f"任务失败: {task.error_message}")
            
            # 等待10秒
            time.sleep(10)
        
        # 检查是否超时
        if not completed:
            self.fail(f"任务执行超时 ({timeout}秒)")
        
        # 验证结果
        self.assertEqual(task.status, 5, "任务状态应为'已完成'")
        self.assertIsNotNone(task.final_video_path, "应生成最终视频文件")
        self.assertTrue(os.path.exists(task.final_video_path), f"视频文件不存在: {task.final_video_path}")
        
        # 验证视频大小
        video_size = os.path.getsize(task.final_video_path)
        self.assertGreater(video_size, 1024*1024, "视频文件过小，可能生成失败")
        
        logging.info(f"任务完成，最终视频: {task.final_video_path}")
        logging.info(f"视频大小: {video_size/1024/1024:.2f} MB")

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