epub_bilingual_translator/archive/v0.09/tests/test_translator.py

"""
翻译器测试
"""

import pytest
import asyncio
from unittest.mock import Mock, AsyncMock, patch

from src.translator import EPUBTranslator
from src.llm_client import OpenRouterClient


class TestEPUBTranslator:
    """EPUB 翻译器测试类"""
    
    @pytest.fixture
    def mock_config(self):
        """模拟配置"""
        return {
            'openrouter': {
                'api_key': 'test_key',
                'base_url': 'https://openrouter.ai/api/v1',
                'models': {
                    'test': 'google/gemini-2.0-flash-exp',
                    'production': 'google/gemini-exp-1206'
                },
                'rate_limits': {
                    'requests_per_minute': 60,
                    'concurrent_requests': 5
                }
            },
            'translation': {
                'chunk_size': 3,
                'max_context_length': 8000,
                'sample_ratio': 0.1,
                'target_language': 'zh-CN',
                'temperature': 0.3,
                'max_tokens': 4000
            },
            'processing': {
                'skip_sections': ['acknowledgments'],
                'include_sections': ['preface', 'chapter'],
                'clean_patterns': ['\\[\\d+\\]'],
                'min_paragraph_length': 20
            },
            'output': {
                'format': 'bilingual',
                'filename_suffix': '_bilingual',
                'preserve_images': True,
                'preserve_css': True,
                'output_dir': 'output'
            },
            'logging': {
                'level': 'INFO',
                'file': 'logs/test.log'
            }
        }
    
    @pytest.fixture
    def mock_translator(self, mock_config):
        """创建模拟翻译器"""
        with patch('src.translator.OpenRouterClient') as mock_client:
            mock_client.return_value.close = AsyncMock()
            translator = EPUBTranslator(mock_config)
            return translator
    
    def test_translator_initialization(self, mock_translator):
        """测试翻译器初始化"""
        assert mock_translator.config is not None
        assert mock_translator.text_processor is not None
        assert mock_translator.console is not None
    
    @pytest.mark.asyncio
    async def test_get_translation_estimate(self, mock_translator):
        """测试翻译估算"""
        # 模拟 EPUBParser
        with patch('src.translator.EPUBParser') as mock_parser:
            mock_parser.return_value.extract_translatable_content.return_value = [
                {
                    'title': 'Test Chapter',
                    'content': '<p>Test paragraph content</p>' * 10,
                    'type': 'chapter'
                }
            ]
            
            # 模拟 text_processor
            mock_translator.text_processor.extract_paragraphs = Mock(return_value=[
                {'text': 'Test paragraph content'} for _ in range(10)
            ])
            
            estimate = await mock_translator.get_translation_estimate('test.epub')
            
            assert 'total_paragraphs' in estimate
            assert 'estimated_tokens' in estimate
            assert 'estimated_time_minutes' in estimate
    
    def test_get_translator_info(self, mock_translator):
        """测试获取翻译器信息"""
        # 模拟 llm_client
        mock_translator.llm_client.get_model_info = Mock(return_value={
            'test_model': 'test_model',
            'production_model': 'prod_model'
        })
        
        info = mock_translator.get_translator_info()
        
        assert 'version' in info
        assert 'llm_models' in info
        assert 'config' in info
        assert info['config']['chunk_size'] == 3


class TestOpenRouterClient:
    """OpenRouter 客户端测试类"""
    
    @pytest.fixture
    def mock_config(self):
        """模拟配置"""
        return {
            'openrouter': {
                'api_key': 'test_key',
                'base_url': 'https://openrouter.ai/api/v1',
                'models': {
                    'test': 'google/gemini-2.0-flash-exp',
                    'production': 'google/gemini-exp-1206'
                },
                'rate_limits': {
                    'requests_per_minute': 60,
                    'concurrent_requests': 5
                }
            },
            'translation': {
                'temperature': 0.3,
                'max_tokens': 4000
            }
        }
    
    def test_client_initialization_invalid_key(self, mock_config):
        """测试无效 API Key"""
        mock_config['openrouter']['api_key'] = 'YOUR_OPENROUTER_API_KEY'
        
        with pytest.raises(ValueError, match="请在配置文件中设置有效的 OpenRouter API Key"):
            OpenRouterClient(mock_config)
    
    @patch('src.llm_client.AsyncOpenAI')
    def test_client_initialization_valid(self, mock_openai, mock_config):
        """测试有效初始化"""
        client = OpenRouterClient(mock_config)
        
        assert client.models['test'] == 'google/gemini-2.0-flash-exp'
        assert client.models['production'] == 'google/gemini-exp-1206'
        mock_openai.assert_called_once()
    
    @patch('src.llm_client.AsyncOpenAI')
    def test_build_translation_prompt(self, mock_openai, mock_config):
        """测试翻译提示词构建"""
        client = OpenRouterClient(mock_config)
        
        prompt = client.build_translation_prompt(
            "Hello world",
            "This is a test book",
            {"technical_terms": {"API": "应用程序接口"}}
        )
        
        assert "Hello world" in prompt
        assert "This is a test book" in prompt
        assert "API -> 应用程序接口" in prompt
    
    @patch('src.llm_client.AsyncOpenAI')
    def test_split_translation_result(self, mock_openai, mock_config):
        """测试翻译结果分割"""
        client = OpenRouterClient(mock_config)
        
        # 测试正常分割
        translation = "第一段翻译\n\n第二段翻译\n\n第三段翻译"
        result = client._split_translation_result(translation, 3)
        
        assert len(result) == 3
        assert result[0] == "第一段翻译"
        assert result[1] == "第二段翻译"
        assert result[2] == "第三段翻译"
        
        # 测试单段落
        single_translation = "单段落翻译"
        result = client._split_translation_result(single_translation, 1)
        
        assert len(result) == 1
        assert result[0] == "单段落翻译"