openclaw-skills/skills/obsidian-kb/test_subtitle_generator.py

#!/usr/bin/env python3
"""
字幕生成器测试脚本

测试各个模块的功能
"""

import os
import sys
import tempfile
from pathlib import Path

# 添加项目路径
sys.path.insert(0, str(Path(__file__).parent))


def test_audio_downloader():
    """测试音频下载器"""
    print("\n" + "="*60)
    print("测试 1: 音频下载器")
    print("="*60)

    from subtitle_generator import AudioDownloader

    # 使用一个短测试视频
    test_url = "https://www.youtube.com/watch?v=dQw4w9WgXcQ"  # Rick Astley - Never Gonna Give You Up

    downloader = AudioDownloader()

    try:
        with tempfile.TemporaryDirectory() as tmpdir:
            output_path = os.path.join(tmpdir, "test_audio.%(ext)s")
            result_path = downloader.download_audio(test_url, output_path)

            if os.path.exists(result_path):
                size_mb = os.path.getsize(result_path) / (1024 * 1024)
                print(f"✅ 音频下载成功")
                print(f"   文件: {result_path}")
                print(f"   大小: {size_mb:.2f} MB")
                return True
            else:
                print(f"❌ 音频文件未找到: {result_path}")
                return False

    except Exception as e:
        print(f"❌ 音频下载失败: {e}")
        import traceback
        traceback.print_exc()
        return False


def test_whisper_transcriber():
    """测试 Whisper 识别器"""
    print("\n" + "="*60)
    print("测试 2: Whisper 识别器")
    print("="*60)

    from subtitle_generator import WhisperTranscriber

    # 创建测试音频（使用 ffmpeg 生成静音）
    test_audio = "/tmp/test_silent.mp3"
    os.system(f"ffmpeg -f lavfi -i anullsrc=r=44100:cl=mono -t 3 -q:a 9 -y {test_audio} 2>/dev/null")

    try:
        transcriber = WhisperTranscriber(model_size='tiny')  # 使用最小的模型
        result = transcriber.transcribe(test_audio)

        print(f"✅ Whisper 识别完成")
        print(f"   识别文本: {result['text'][:50]}...")
        print(f"   分段数: {len(result.get('segments', []))}")

        return True

    except Exception as e:
        print(f"❌ Whisper 识别失败: {e}")
        import traceback
        traceback.print_exc()
        return False


def test_subtitle_formatter():
    """测试 SRT 格式化器"""
    print("\n" + "="*60)
    print("测试 3: SRT 格式化器")
    print("="*60)

    from subtitle_generator import SubtitleFormatter

    # 测试数据
    test_segments = [
        {'start': 0.0, 'end': 2.5, 'text': '这是一段测试字幕'},
        {'start': 3.0, 'end': 5.7, 'text': '这是第二段字幕'},
        {'start': 6.0, 'end': 8.2, 'text': 'This is the third subtitle'},
    ]

    try:
        srt_text = SubtitleFormatter.to_srt(test_segments)

        print(f"✅ SRT 格式化成功")
        print(f"   生成内容预览:")
        print("   " + "-" * 40)
        for line in srt_text.split('\n')[:10]:
            print(f"   {line}")
        print("   " + "-" * 40)

        # 验证格式
        assert '00:00:00,000 --> 00:00:02,500' in srt_text
        assert '这是一段测试字幕' in srt_text

        return True

    except Exception as e:
        print(f"❌ SRT 格式化失败: {e}")
        import traceback
        traceback.print_exc()
        return False


def test_subtitle_cache():
    """测试缓存管理器"""
    print("\n" + "="*60)
    print("测试 4: 缓存管理器")
    print("="*60)

    from subtitle_generator import SubtitleCache

    test_url = "https://example.com/video/123"
    test_subtitle = """1
00:00:00,000 --> 00:00:02,500
测试字幕内容
"""

    try:
        cache = SubtitleCache(cache_dir='/tmp/test_subtitle_cache')

        # 测试写入
        cache_path = cache.set(test_url, test_subtitle)
        print(f"✅ 缓存写入成功: {cache_path}")

        # 测试读取
        cached_content = cache.get(test_url)
        assert cached_content == test_subtitle
        print(f"✅ 缓存读取成功")

        # 测试存在性
        assert cache.exists(test_url)
        print(f"✅ 缓存存在性检查成功")

        # 测试缓存键
        cache_key = cache.get_cache_key(test_url)
        print(f"   缓存键: {cache_key}")

        return True

    except Exception as e:
        print(f"❌ 缓存测试失败: {e}")
        import traceback
        traceback.print_exc()
        return False


def test_full_workflow():
    """测试完整工作流程"""
    print("\n" + "="*60)
    print("测试 5: 完整工作流程")
    print("="*60)

    from subtitle_generator import SubtitleGenerator

    test_url = "https://www.youtube.com/watch?v=dQw4w9WgXcQ"

    try:
        generator = SubtitleGenerator(model_size='tiny')

        print(f"开始生成字幕...")
        srt_text = generator.generate(
            url=test_url,
            content_id='test001',
            language='en',
            use_cache=False
        )

        print(f"✅ 字幕生成完成")
        print(f"   总字符数: {len(srt_text)}")
        print(f"   预览（前300字符）:")
        print("   " + "-" * 40)
        print("   " + srt_text[:300])
        print("   " + "-" * 40)

        # 验证缓存
        print(f"\n验证缓存...")
        cached = generator.generate(
            url=test_url,
            content_id='test002',
            use_cache=True
        )
        assert cached == srt_text
        print(f"✅ 缓存验证成功")

        return True

    except Exception as e:
        print(f"❌ 完整工作流程失败: {e}")
        import traceback
        traceback.print_exc()
        return False


def main():
    """运行所有测试"""
    print("\n" + "="*60)
    print("字幕生成器测试套件")
    print("="*60)

    results = {}

    # 运行各个测试
    results['audio_downloader'] = test_audio_downloader()
    results['whisper_transcriber'] = test_whisper_transcriber()
    results['subtitle_formatter'] = test_subtitle_formatter()
    results['subtitle_cache'] = test_subtitle_cache()
    results['full_workflow'] = test_full_workflow()

    # 输出汇总
    print("\n" + "="*60)
    print("测试结果汇总")
    print("="*60)

    for test_name, passed in results.items():
        status = "✅ 通过" if passed else "❌ 失败"
        print(f"{test_name:25s} {status}")

    total = len(results)
    passed = sum(results.values())

    print("\n" + "="*60)
    print(f"总计: {passed}/{total} 测试通过")
    print("="*60)

    return all(results.values())


if __name__ == '__main__':
    success = main()
    sys.exit(0 if success else 1)