tests/test_ai_model.py

# tests/test_ai_model.py
import pytest
from unittest.mock import Mock, patch
from modules.ai_model import AIModel

class TestAIModel:
    
    @pytest.fixture
    def mock_ai_model(self):
        """创建模拟的AI模型"""
        with patch('modules.ai_model.Gemma3nForConditionalGeneration') as mock_model, \
             patch('modules.ai_model.AutoProcessor') as mock_processor:
            
            mock_model.from_pretrained.return_value = Mock()
            mock_processor.from_pretrained.return_value = Mock()
            
            ai_model = AIModel()
            yield ai_model
    
    def test_detect_input_type_text(self, mock_ai_model):
        """测试文本输入检测"""
        assert mock_ai_model.detect_input_type("我想去巴黎") == "text"
    
    def test_detect_input_type_image_url(self, mock_ai_model):
        """测试图片URL检测"""
        assert mock_ai_model.detect_input_type("https://example.com/image.jpg") == "image"
    
    def test_detect_input_type_image_path(self, mock_ai_model):
        """测试图片路径检测"""
        assert mock_ai_model.detect_input_type("./images/paris.png") == "image"
    
    def test_detect_input_type_audio(self, mock_ai_model):
        """测试音频文件检测"""
        assert mock_ai_model.detect_input_type("audio.mp3") == "audio"
    
    def test_is_available(self, mock_ai_model):
        """测试模型可用性检查"""
        assert mock_ai_model.is_available() == True
    
    @patch('modules.ai_model.Image')
    @patch('modules.ai_model.requests')
    def test_format_input_image_url(self, mock_requests, mock_image, mock_ai_model):
        """测试图片URL格式化"""
        mock_response = Mock()
        mock_response.content = b"fake_image_data"
        mock_requests.get.return_value = mock_response
        mock_image.open.return_value.convert.return_value = "processed_image"
        
        input_type, formatted_data, processed_text = mock_ai_model.format_input(
            "image", 
            "https://example.com/image.jpg"
        )
        
        assert input_type == "image"
        assert formatted_data == "processed_image"
        assert "请描述这张图片" in processed_text