tests/unit/test_model_config.py

"""
Unit tests for model configuration components.
"""

import pytest
import torch
from unittest.mock import patch, MagicMock
import os
import sys
from pathlib import Path

# Add project root to path
project_root = Path(__file__).parent.parent.parent
sys.path.insert(0, str(project_root))

from app.config.model_config import (
    ModelConfig, 
    EnvironmentDetector, 
    DependencyValidator
)


class TestModelConfig:
    """Test ModelConfig dataclass."""
    
    def test_model_config_creation(self):
        """Test ModelConfig creation with valid parameters."""
        config = ModelConfig(
            model_id="test/model",
            revision="main",
            dtype=torch.float32,
            device_map="cpu",
            attn_implementation="eager",
            low_cpu_mem_usage=True,
            trust_remote_code=True
        )
        
        assert config.model_id == "test/model"
        assert config.revision == "main"
        assert config.dtype == torch.float32
        assert config.device_map == "cpu"
        assert config.attn_implementation == "eager"
        assert config.low_cpu_mem_usage is True
        assert config.trust_remote_code is True
    
    @patch('torch.cuda.is_available')
    def test_is_gpu_available_true(self, mock_cuda):
        """Test GPU availability detection when CUDA is available."""
        mock_cuda.return_value = True
        
        config = ModelConfig(
            model_id="test/model",
            revision=None,
            dtype=torch.float32,
            device_map="auto",
            attn_implementation="sdpa",
            low_cpu_mem_usage=False,
            trust_remote_code=True
        )
        
        assert config.is_gpu_available is True
    
    @patch('torch.cuda.is_available')
    def test_is_gpu_available_false(self, mock_cuda):
        """Test GPU availability detection when CUDA is not available."""
        mock_cuda.return_value = False
        
        config = ModelConfig(
            model_id="test/model",
            revision=None,
            dtype=torch.float32,
            device_map="cpu",
            attn_implementation="eager",
            low_cpu_mem_usage=True,
            trust_remote_code=True
        )
        
        assert config.is_gpu_available is False
    
    @patch('torch.cuda.is_available')
    @patch('torch.cuda.device_count')
    @patch('torch.cuda.current_device')
    @patch('torch.cuda.get_device_name')
    @patch('torch.cuda.memory_allocated')
    @patch('torch.cuda.memory_reserved')
    def test_device_info_gpu(self, mock_mem_reserved, mock_mem_allocated, 
                           mock_device_name, mock_current_device, 
                           mock_device_count, mock_cuda):
        """Test device info when GPU is available."""
        mock_cuda.return_value = True
        mock_device_count.return_value = 1
        mock_current_device.return_value = 0
        mock_device_name.return_value = "Test GPU"
        mock_mem_allocated.return_value = 1024
        mock_mem_reserved.return_value = 2048
        
        config = ModelConfig(
            model_id="test/model",
            revision=None,
            dtype=torch.float32,
            device_map="auto",
            attn_implementation="sdpa",
            low_cpu_mem_usage=False,
            trust_remote_code=True
        )
        
        device_info = config.device_info
        
        assert device_info["cuda_available"] is True
        assert device_info["device_count"] == 1
        assert device_info["current_device"] == 0
        assert device_info["device_name"] == "Test GPU"
        assert device_info["memory_allocated"] == 1024
        assert device_info["memory_reserved"] == 2048
    
    @patch('torch.cuda.is_available')
    def test_device_info_cpu(self, mock_cuda):
        """Test device info when only CPU is available."""
        mock_cuda.return_value = False
        
        config = ModelConfig(
            model_id="test/model",
            revision=None,
            dtype=torch.float32,
            device_map="cpu",
            attn_implementation="eager",
            low_cpu_mem_usage=True,
            trust_remote_code=True
        )
        
        device_info = config.device_info
        
        assert device_info["cuda_available"] is False
        assert device_info["device_count"] == 0
        assert device_info["current_device"] is None


class TestEnvironmentDetector:
    """Test EnvironmentDetector class."""
    
    @patch('torch.cuda.is_available')
    @patch('torch.__version__', '2.0.0')
    def test_detect_environment_cpu(self, mock_cuda):
        """Test environment detection on CPU-only system."""
        mock_cuda.return_value = False
        
        with patch('importlib.import_module') as mock_import:
            # Mock successful einops import
            mock_einops = MagicMock()
            mock_einops.__version__ = "0.7.0"
            
            def import_side_effect(name):
                if name == "einops":
                    return mock_einops
                elif name == "flash_attn":
                    raise ImportError("No module named 'flash_attn'")
                return MagicMock()
            
            mock_import.side_effect = import_side_effect
            
            env_info = EnvironmentDetector.detect_environment()
            
            assert env_info["cuda_available"] is False
            assert env_info["cuda_version"] is None
            assert env_info["torch_version"] == "2.0.0"
    
    @patch('torch.cuda.is_available')
    @patch('torch.version.cuda', '11.8')
    def test_detect_environment_gpu(self, mock_cuda):
        """Test environment detection on GPU system."""
        mock_cuda.return_value = True
        
        env_info = EnvironmentDetector.detect_environment()
        
        assert env_info["cuda_available"] is True
        assert env_info["cuda_version"] == "11.8"
    
    @patch('torch.cuda.is_available')
    @patch.dict(os.environ, {}, clear=True)
    def test_create_model_config_cpu_defaults(self, mock_cuda):
        """Test model config creation with CPU defaults."""
        mock_cuda.return_value = False
        
        config = EnvironmentDetector.create_model_config()
        
        assert config.model_id == "microsoft/Phi-3.5-MoE-instruct"
        assert config.revision is None
        assert config.dtype == torch.float32
        assert config.device_map == "cpu"
        assert config.attn_implementation == "eager"
        assert config.low_cpu_mem_usage is True
        assert config.trust_remote_code is True
    
    @patch('torch.cuda.is_available')
    @patch.dict(os.environ, {}, clear=True)
    def test_create_model_config_gpu_defaults(self, mock_cuda):
        """Test model config creation with GPU defaults."""
        mock_cuda.return_value = True
        
        config = EnvironmentDetector.create_model_config()
        
        assert config.model_id == "microsoft/Phi-3.5-MoE-instruct"
        assert config.revision is None
        assert config.dtype == torch.bfloat16
        assert config.device_map == "auto"
        assert config.attn_implementation == "sdpa"
        assert config.low_cpu_mem_usage is False
        assert config.trust_remote_code is True
    
    @patch('torch.cuda.is_available')
    @patch.dict(os.environ, {
        'HF_MODEL_ID': 'custom/model',
        'HF_REVISION': 'abc123'
    })
    def test_create_model_config_with_env_vars(self, mock_cuda):
        """Test model config creation with environment variables."""
        mock_cuda.return_value = False
        
        config = EnvironmentDetector.create_model_config()
        
        assert config.model_id == "custom/model"
        assert config.revision == "abc123"
    
    def test_create_model_config_with_parameters(self):
        """Test model config creation with explicit parameters."""
        config = EnvironmentDetector.create_model_config(
            model_id="explicit/model",
            revision="explicit123"
        )
        
        assert config.model_id == "explicit/model"
        assert config.revision == "explicit123"


class TestDependencyValidator:
    """Test DependencyValidator class."""
    
    def test_required_packages_list(self):
        """Test that required packages list is properly defined."""
        required = DependencyValidator.REQUIRED_PACKAGES
        
        assert "transformers" in required
        assert "accelerate" in required
        assert "einops" in required
        assert "huggingface_hub" in required
        assert "gradio" in required
        assert "torch" in required
    
    def test_optional_packages_list(self):
        """Test that optional packages list is properly defined."""
        optional = DependencyValidator.OPTIONAL_PACKAGES
        
        assert "flash_attn" in optional
    
    @patch('importlib.import_module')
    def test_validate_dependencies_all_available(self, mock_import):
        """Test dependency validation when all packages are available."""
        mock_import.return_value = MagicMock()
        
        results = DependencyValidator.validate_dependencies()
        
        for package in DependencyValidator.REQUIRED_PACKAGES:
            assert results[package] is True
        
        for package in DependencyValidator.OPTIONAL_PACKAGES:
            assert results[package] is True
    
    @patch('importlib.import_module')
    def test_validate_dependencies_some_missing(self, mock_import):
        """Test dependency validation when some packages are missing."""
        def import_side_effect(name):
            if name == "flash_attn":
                raise ImportError(f"No module named '{name}'")
            return MagicMock()
        
        mock_import.side_effect = import_side_effect
        
        results = DependencyValidator.validate_dependencies()
        
        # Required packages should be available
        for package in DependencyValidator.REQUIRED_PACKAGES:
            assert results[package] is True
        
        # flash_attn should be missing
        assert results["flash_attn"] is False
    
    @patch('importlib.import_module')
    def test_get_missing_required_packages(self, mock_import):
        """Test getting list of missing required packages."""
        def import_side_effect(name):
            if name in ["transformers", "einops"]:
                raise ImportError(f"No module named '{name}'")
            return MagicMock()
        
        mock_import.side_effect = import_side_effect
        
        missing = DependencyValidator.get_missing_required_packages()
        
        assert "transformers" in missing
        assert "einops" in missing
        assert "accelerate" not in missing  # Should be available
    
    @patch('importlib.import_module')
    def test_is_environment_ready_true(self, mock_import):
        """Test environment readiness when all required packages are available."""
        mock_import.return_value = MagicMock()
        
        assert DependencyValidator.is_environment_ready() is True
    
    @patch('importlib.import_module')
    def test_is_environment_ready_false(self, mock_import):
        """Test environment readiness when required packages are missing."""
        def import_side_effect(name):
            if name == "einops":
                raise ImportError(f"No module named '{name}'")
            return MagicMock()
        
        mock_import.side_effect = import_side_effect
        
        assert DependencyValidator.is_environment_ready() is False


if __name__ == "__main__":
    pytest.main([__file__])