Create tests/test_models.py

tests/test_models.py

@@ -1,349 +1,376 @@
 """
-Tests for the processing pipeline.
+Tests for model management functionality.
 """
 
 import pytest
-import numpy as np
-import cv2
-from unittest.mock import Mock, patch, MagicMock
+import tempfile
 from pathlib import Path
+from unittest.mock import Mock, patch, MagicMock
+import json
 
-from api.pipeline import (
-    ProcessingPipeline,
-    PipelineConfig,
-    PipelineResult,
-    ProcessingMode,
-    PipelineStage
+from models import (
+    ModelRegistry,
+    ModelInfo,
+    ModelStatus,
+    ModelTask,
+    ModelFramework,
+    ModelDownloader,
+    ModelLoader,
+    ModelOptimizer
 )
 
 
-class TestPipelineConfig:
-    """Test pipeline configuration."""
+class TestModelRegistry:
+    """Test model registry functionality."""
+    
+    @pytest.fixture
+    def registry(self):
+        """Create a test registry."""
+        temp_dir = tempfile.mkdtemp()
+        return ModelRegistry(models_dir=Path(temp_dir))
     
-    def test_default_config(self):
-        """Test default configuration values."""
-        config = PipelineConfig()
-        assert config.mode == ProcessingMode.PHOTO
-        assert config.quality_preset == "high"
-        assert config.use_gpu == True
-        assert config.enable_cache == True
+    def test_registry_initialization(self, registry):
+        """Test registry initialization."""
+        assert registry is not None
+        assert len(registry.models) > 0  # Should have default models
+        assert registry.models_dir.exists()
     
-    def test_custom_config(self):
-        """Test custom configuration."""
-        config = PipelineConfig(
-            mode=ProcessingMode.VIDEO,
-            quality_preset="ultra",
-            use_gpu=False,
-            batch_size=4
+    def test_register_model(self, registry):
+        """Test registering a new model."""
+        model = ModelInfo(
+            model_id="test-model",
+            name="Test Model",
+            version="1.0",
+            task=ModelTask.SEGMENTATION,
+            framework=ModelFramework.PYTORCH,
+            url="http://example.com/model.pth",
+            filename="test.pth",
+            file_size=1000000
         )
-        assert config.mode == ProcessingMode.VIDEO
-        assert config.quality_preset == "ultra"
-        assert config.use_gpu == False
-        assert config.batch_size == 4
-
-
-class TestProcessingPipeline:
-    """Test the main processing pipeline."""
+        
+        success = registry.register_model(model)
+        assert success == True
+        assert "test-model" in registry.models
     
-    @pytest.fixture
-    def mock_pipeline(self, pipeline_config):
-        """Create a pipeline with mocked components."""
-        with patch('api.pipeline.ModelFactory') as mock_factory:
-            with patch('api.pipeline.DeviceManager') as mock_device:
-                mock_device.return_value.get_device.return_value = 'cpu'
-                mock_factory.return_value.load_model.return_value = Mock()
-                
-                pipeline = ProcessingPipeline(pipeline_config)
-                return pipeline
+    def test_get_model(self, registry):
+        """Test getting a model by ID."""
+        model = registry.get_model("rmbg-1.4")
+        assert model is not None
+        assert model.model_id == "rmbg-1.4"
+        assert model.task == ModelTask.SEGMENTATION
     
-    def test_pipeline_initialization(self, mock_pipeline):
-        """Test pipeline initialization."""
-        assert mock_pipeline is not None
-        assert mock_pipeline.config is not None
-        assert mock_pipeline.current_stage == PipelineStage.INITIALIZATION
+    def test_list_models_by_task(self, registry):
+        """Test listing models by task."""
+        segmentation_models = registry.list_models(task=ModelTask.SEGMENTATION)
+        assert len(segmentation_models) > 0
+        assert all(m.task == ModelTask.SEGMENTATION for m in segmentation_models)
     
-    def test_process_image_success(self, mock_pipeline, sample_image, sample_background):
-        """Test successful image processing."""
-        # Mock the processing methods
-        mock_pipeline._segment_image = Mock(return_value=np.ones((512, 512), dtype=np.uint8) * 255)
-        mock_pipeline.alpha_matting.process = Mock(return_value={
-            'alpha': np.ones((512, 512), dtype=np.float32),
-            'confidence': 0.95
-        })
-        
-        result = mock_pipeline.process_image(sample_image, sample_background)
-        
-        assert result is not None
-        assert isinstance(result, PipelineResult)
-        assert result.success == True
-        assert result.output_image is not None
+    def test_list_models_by_framework(self, registry):
+        """Test listing models by framework."""
+        pytorch_models = registry.list_models(framework=ModelFramework.PYTORCH)
+        onnx_models = registry.list_models(framework=ModelFramework.ONNX)
+        
+        assert all(m.framework == ModelFramework.PYTORCH for m in pytorch_models)
+        assert all(m.framework == ModelFramework.ONNX for m in onnx_models)
     
-    def test_process_image_with_effects(self, mock_pipeline, sample_image):
-        """Test image processing with effects."""
-        mock_pipeline.config.apply_effects = ['bokeh', 'vignette']
+    def test_get_best_model(self, registry):
+        """Test getting best model for a task."""
+        # Best for accuracy
+        best_accuracy = registry.get_best_model(
+            ModelTask.SEGMENTATION,
+            prefer_speed=False
+        )
+        assert best_accuracy is not None
         
-        # Mock processing
-        mock_pipeline._segment_image = Mock(return_value=np.ones((512, 512), dtype=np.uint8) * 255)
-        mock_pipeline.alpha_matting.process = Mock(return_value={
-            'alpha': np.ones((512, 512), dtype=np.float32),
-            'confidence': 0.95
-        })
+        # Best for speed
+        best_speed = registry.get_best_model(
+            ModelTask.SEGMENTATION,
+            prefer_speed=True
+        )
+        assert best_speed is not None
+    
+    def test_update_model_usage(self, registry):
+        """Test updating model usage statistics."""
+        model_id = "rmbg-1.4"
+        initial_count = registry.models[model_id].use_count
         
-        result = mock_pipeline.process_image(sample_image, None)
+        registry.update_model_usage(model_id)
         
-        assert result is not None
-        assert result.success == True
+        assert registry.models[model_id].use_count == initial_count + 1
+        assert registry.models[model_id].last_used is not None
+    
+    def test_get_total_size(self, registry):
+        """Test calculating total model size."""
+        total_size = registry.get_total_size()
+        assert total_size > 0
+        
+        # Size of available models should be 0 initially
+        available_size = registry.get_total_size(status=ModelStatus.AVAILABLE)
+        assert available_size == 0
     
-    def test_process_image_failure(self, mock_pipeline, sample_image):
-        """Test image processing failure handling."""
-        # Mock segmentation to fail
-        mock_pipeline._segment_image = Mock(side_effect=Exception("Segmentation failed"))
+    def test_export_registry(self, registry, temp_dir):
+        """Test exporting registry to file."""
+        export_path = temp_dir / "registry_export.json"
+        registry.export_registry(export_path)
         
-        result = mock_pipeline.process_image(sample_image, None)
+        assert export_path.exists()
         
-        assert result is not None
-        assert result.success == False
-        assert len(result.errors) > 0
+        with open(export_path) as f:
+            data = json.load(f)
+            assert "models" in data
+            assert len(data["models"]) > 0
+
+
+class TestModelDownloader:
+    """Test model downloading functionality."""
     
-    @pytest.mark.parametrize("quality", ["low", "medium", "high", "ultra"])
-    def test_quality_presets(self, mock_pipeline, sample_image, quality):
-        """Test different quality presets."""
-        mock_pipeline.config.quality_preset = quality
-        
-        # Mock processing
-        mock_pipeline._segment_image = Mock(return_value=np.ones((512, 512), dtype=np.uint8) * 255)
-        mock_pipeline.alpha_matting.process = Mock(return_value={
-            'alpha': np.ones((512, 512), dtype=np.float32),
-            'confidence': 0.95
-        })
-        
-        result = mock_pipeline.process_image(sample_image, None)
-        
-        assert result is not None
-        assert result.success == True
+    @pytest.fixture
+    def downloader(self, mock_registry):
+        """Create a test downloader."""
+        return ModelDownloader(mock_registry)
     
-    def test_batch_processing(self, mock_pipeline, sample_image):
-        """Test batch processing of multiple images."""
-        images = [sample_image] * 3
-        
-        # Mock processing
-        mock_pipeline.process_image = Mock(return_value=PipelineResult(
-            success=True,
-            output_image=sample_image,
-            quality_score=0.9
-        ))
+    @patch('requests.get')
+    def test_download_model(self, mock_get, downloader):
+        """Test downloading a model."""
+        # Mock HTTP response
+        mock_response = MagicMock()
+        mock_response.headers = {'content-length': '1000000'}
+        mock_response.iter_content = MagicMock(
+            return_value=[b'data' * 1000]
+        )
+        mock_response.raise_for_status = MagicMock()
+        mock_get.return_value = mock_response
         
-        results = mock_pipeline.process_batch(images)
+        # Test download
+        success = downloader.download_model("test-model", force=True)
         
-        assert len(results) == 3
-        assert all(r.success for r in results)
+        assert mock_get.called
+        # Note: Full download test would require more mocking
     
-    def test_progress_callback(self, mock_pipeline, sample_image):
-        """Test progress callback functionality."""
+    def test_download_progress_tracking(self, downloader):
+        """Test download progress tracking."""
         progress_values = []
         
-        def progress_callback(value, message):
-            progress_values.append(value)
+        def progress_callback(progress):
+            progress_values.append(progress.progress)
         
-        mock_pipeline.config.progress_callback = progress_callback
+        # Start a download (will fail but we can test progress initialization)
+        with patch.object(downloader, '_download_model_task', return_value=True):
+            downloader.download_model(
+                "test-model",
+                progress_callback=progress_callback
+            )
         
-        # Mock processing
-        mock_pipeline._segment_image = Mock(return_value=np.ones((512, 512), dtype=np.uint8) * 255)
-        mock_pipeline.alpha_matting.process = Mock(return_value={
-            'alpha': np.ones((512, 512), dtype=np.float32),
-            'confidence': 0.95
-        })
+        assert "test-model" in downloader.downloads
+    
+    def test_cancel_download(self, downloader):
+        """Test cancelling a download."""
+        # Start a mock download
+        downloader.downloads["test-model"] = Mock()
+        downloader._stop_events["test-model"] = Mock()
         
-        result = mock_pipeline.process_image(sample_image, None)
+        success = downloader.cancel_download("test-model")
         
-        assert len(progress_values) > 0
-        assert 0.0 <= max(progress_values) <= 1.0
+        assert success == True
+        assert downloader._stop_events["test-model"].set.called
     
-    def test_cache_functionality(self, mock_pipeline, sample_image):
-        """Test caching functionality."""
-        mock_pipeline.config.enable_cache = True
-        
-        # Mock processing
-        mock_pipeline._segment_image = Mock(return_value=np.ones((512, 512), dtype=np.uint8) * 255)
-        mock_pipeline.alpha_matting.process = Mock(return_value={
-            'alpha': np.ones((512, 512), dtype=np.float32),
-            'confidence': 0.95
-        })
-        
-        # First call
-        result1 = mock_pipeline.process_image(sample_image, None)
-        
-        # Second call (should use cache)
-        result2 = mock_pipeline.process_image(sample_image, None)
-        
-        assert result1.success == result2.success
-        # Verify segmentation was only called once (cache hit on second call)
-        assert mock_pipeline._segment_image.call_count == 1
+    def test_download_with_resume(self, downloader, temp_dir):
+        """Test download with resume support."""
+        # Create a partial file
+        partial_file = temp_dir / "test.pth.part"
+        partial_file.write_bytes(b"partial_data")
+        
+        # Mock download would check for partial file
+        assert partial_file.exists()
+        assert partial_file.stat().st_size > 0
+
+
+class TestModelLoader:
+    """Test model loading functionality."""
+    
+    @pytest.fixture
+    def loader(self, mock_registry):
+        """Create a test loader."""
+        return ModelLoader(mock_registry, device='cpu')
     
-    def test_memory_management(self, mock_pipeline):
-        """Test memory management and cleanup."""
-        initial_cache_size = len(mock_pipeline.cache)
+    def test_loader_initialization(self, loader):
+        """Test loader initialization."""
+        assert loader is not None
+        assert loader.device == 'cpu'
+        assert loader.max_memory_bytes > 0
+    
+    @patch('torch.load')
+    def test_load_pytorch_model(self, mock_torch_load, loader):
+        """Test loading a PyTorch model."""
+        mock_model = MagicMock()
+        mock_torch_load.return_value = mock_model
+        
+        # Mock model info
+        model_info = ModelInfo(
+            model_id="test-pytorch",
+            name="Test PyTorch Model",
+            version="1.0",
+            task=ModelTask.SEGMENTATION,
+            framework=ModelFramework.PYTORCH,
+            url="",
+            filename="model.pth",
+            local_path="/tmp/model.pth",
+            status=ModelStatus.AVAILABLE
+        )
         
-        # Process multiple images to fill cache
-        for i in range(10):
-            image = np.random.randint(0, 255, (512, 512, 3), dtype=np.uint8)
-            mock_pipeline.cache[f"test_{i}"] = PipelineResult(success=True)
+        loader.registry.get_model = Mock(return_value=model_info)
         
-        # Clear cache
-        mock_pipeline.clear_cache()
+        with patch.object(Path, 'exists', return_value=True):
+            loaded = loader.load_model("test-pytorch")
         
-        assert len(mock_pipeline.cache) == 0
+        # Note: Full test would require more setup
+        assert mock_torch_load.called
     
-    def test_statistics_tracking(self, mock_pipeline, sample_image):
-        """Test statistics tracking."""
-        # Mock processing
-        mock_pipeline._segment_image = Mock(return_value=np.ones((512, 512), dtype=np.uint8) * 255)
-        mock_pipeline.alpha_matting.process = Mock(return_value={
-            'alpha': np.ones((512, 512), dtype=np.float32),
-            'confidence': 0.95
-        })
-        
-        # Process image
-        result = mock_pipeline.process_image(sample_image, None)
-        
-        # Get statistics
-        stats = mock_pipeline.get_statistics()
-        
-        assert 'total_processed' in stats
-        assert stats['total_processed'] > 0
-        assert 'avg_time' in stats
-
-
-class TestPipelineIntegration:
-    """Integration tests for the pipeline."""
+    def test_memory_management(self, loader):
+        """Test memory management during model loading."""
+        # Add mock models to loaded cache
+        for i in range(5):
+            loader.loaded_models[f"model_{i}"] = Mock(
+                memory_usage=100 * 1024 * 1024  # 100MB each
+            )
+        
+        loader.current_memory_usage = 500 * 1024 * 1024  # 500MB
+        
+        # Free memory
+        loader._free_memory(200 * 1024 * 1024)  # Need 200MB
+        
+        # Should have freed at least 2 models
+        assert len(loader.loaded_models) < 5
     
-    @pytest.mark.integration
-    @pytest.mark.slow
-    def test_end_to_end_processing(self, sample_image, sample_background, temp_dir):
-        """Test end-to-end processing pipeline."""
-        config = PipelineConfig(
-            use_gpu=False,
-            quality_preset="medium",
-            enable_cache=False
+    def test_unload_model(self, loader):
+        """Test unloading a model."""
+        # Add a mock model
+        loader.loaded_models["test"] = Mock(
+            model=Mock(),
+            memory_usage=100 * 1024 * 1024
         )
+        loader.current_memory_usage = 100 * 1024 * 1024
         
-        # Create pipeline (will use real components if available)
-        try:
-            pipeline = ProcessingPipeline(config)
-        except Exception:
-            pytest.skip("Models not available for integration test")
-        
-        # Process image
-        result = pipeline.process_image(sample_image, sample_background)
-        
-        if result.success:
-            assert result.output_image is not None
-            assert result.output_image.shape == sample_image.shape
-            assert result.quality_score > 0
-            
-            # Save output
-            output_path = temp_dir / "test_output.png"
-            cv2.imwrite(str(output_path), result.output_image)
-            assert output_path.exists()
+        success = loader.unload_model("test")
+        
+        assert success == True
+        assert "test" not in loader.loaded_models
+        assert loader.current_memory_usage == 0
     
-    @pytest.mark.integration
-    @pytest.mark.slow
-    def test_video_frame_processing(self, sample_video, temp_dir):
-        """Test processing video frames."""
-        config = PipelineConfig(
-            mode=ProcessingMode.VIDEO,
-            use_gpu=False,
-            quality_preset="low"
+    def test_get_memory_usage(self, loader):
+        """Test getting memory usage statistics."""
+        # Add mock models
+        loader.loaded_models["model1"] = Mock(memory_usage=100 * 1024 * 1024)
+        loader.loaded_models["model2"] = Mock(memory_usage=200 * 1024 * 1024)
+        loader.current_memory_usage = 300 * 1024 * 1024
+        
+        usage = loader.get_memory_usage()
+        
+        assert usage["current_usage_mb"] == 300
+        assert usage["loaded_models"] == 2
+        assert "model1" in usage["models"]
+        assert "model2" in usage["models"]
+
+
+class TestModelOptimizer:
+    """Test model optimization functionality."""
+    
+    @pytest.fixture
+    def optimizer(self, mock_registry):
+        """Create a test optimizer."""
+        loader = ModelLoader(mock_registry, device='cpu')
+        return ModelOptimizer(loader)
+    
+    @patch('torch.quantization.quantize_dynamic')
+    def test_quantize_pytorch_model(self, mock_quantize, optimizer):
+        """Test PyTorch model quantization."""
+        # Create mock model
+        mock_model = MagicMock()
+        mock_quantize.return_value = mock_model
+        
+        loaded = Mock(
+            model_id="test",
+            model=mock_model,
+            framework=ModelFramework.PYTORCH,
+            metadata={'input_size': (1, 3, 512, 512)}
         )
         
-        try:
-            pipeline = ProcessingPipeline(config)
-        except Exception:
-            pytest.skip("Models not available for integration test")
+        with patch.object(optimizer, '_get_model_size', return_value=1000000):
+            with patch.object(optimizer, '_benchmark_model', return_value=0.1):
+                result = optimizer._quantize_pytorch(
+                    loaded,
+                    Path("/tmp"),
+                    "dynamic"
+                )
         
-        # Open video
-        cap = cv2.VideoCapture(sample_video)
-        processed_frames = []
+        assert mock_quantize.called
+        # Note: Full test would require more setup
+    
+    def test_optimization_result(self, optimizer):
+        """Test optimization result structure."""
+        from models.optimizer import OptimizationResult
+        
+        result = OptimizationResult(
+            original_size_mb=100,
+            optimized_size_mb=25,
+            compression_ratio=4.0,
+            original_speed_ms=100,
+            optimized_speed_ms=50,
+            speedup=2.0,
+            accuracy_loss=0.01,
+            optimization_time=10.0,
+            output_path="/tmp/optimized.pth"
+        )
         
-        # Process first 5 frames
-        for i in range(5):
-            ret, frame = cap.read()
-            if not ret:
-                break
-            
-            result = pipeline.process_image(frame, None)
-            if result.success:
-                processed_frames.append(result.output_image)
-        
-        cap.release()
-        
-        assert len(processed_frames) > 0
-        
-        # Save as video
-        if processed_frames:
-            output_path = temp_dir / "test_video_out.mp4"
-            fourcc = cv2.VideoWriter_fourcc(*'mp4v')
-            out = cv2.VideoWriter(str(output_path), fourcc, 30.0, 
-                                (processed_frames[0].shape[1], processed_frames[0].shape[0]))
-            
-            for frame in processed_frames:
-                out.write(frame)
-            
-            out.release()
-            assert output_path.exists()
+        assert result.compression_ratio == 4.0
+        assert result.speedup == 2.0
+        assert result.accuracy_loss == 0.01
 
 
-class TestPipelinePerformance:
-    """Performance tests for the pipeline."""
-    
-    @pytest.mark.slow
-    def test_processing_speed(self, mock_pipeline, sample_image, performance_timer):
-        """Test processing speed."""
-        # Mock processing
-        mock_pipeline._segment_image = Mock(return_value=np.ones((512, 512), dtype=np.uint8) * 255)
-        mock_pipeline.alpha_matting.process = Mock(return_value={
-            'alpha': np.ones((512, 512), dtype=np.float32),
-            'confidence': 0.95
-        })
-        
-        with performance_timer as timer:
-            result = mock_pipeline.process_image(sample_image, None)
-        
-        assert result.success == True
-        assert timer.elapsed < 1.0  # Should process in under 1 second
+class TestModelIntegration:
+    """Integration tests for model management."""
     
+    @pytest.mark.integration
     @pytest.mark.slow
-    def test_batch_processing_speed(self, mock_pipeline, sample_image, performance_timer):
-        """Test batch processing speed."""
-        images = [sample_image] * 10
-        
-        # Mock processing
-        mock_pipeline.process_image = Mock(return_value=PipelineResult(
-            success=True,
-            output_image=sample_image,
-            quality_score=0.9
-        ))
-        
-        with performance_timer as timer:
-            results = mock_pipeline.process_batch(images)
-        
-        assert len(results) == 10
-        assert timer.elapsed < 5.0  # Should process 10 images in under 5 seconds
-    
-    def test_memory_usage(self, mock_pipeline, sample_image):
-        """Test memory usage during processing."""
-        import psutil
-        import os
-        
-        process = psutil.Process(os.getpid())
-        initial_memory = process.memory_info().rss / 1024 / 1024  # MB
+    def test_model_registry_persistence(self, temp_dir):
+        """Test registry persistence across instances."""
+        # Create registry and add model
+        registry1 = ModelRegistry(models_dir=temp_dir)
+        
+        test_model = ModelInfo(
+            model_id="persistence-test",
+            name="Persistence Test",
+            version="1.0",
+            task=ModelTask.SEGMENTATION,
+            framework=ModelFramework.PYTORCH,
+            url="http://example.com/model.pth",
+            filename="persist.pth"
+        )
         
-        # Process multiple images
-        for _ in range(10):
-            mock_pipeline.process_image(sample_image, None)
+        registry1.register_model(test_model)
         
-        final_memory = process.memory_info().rss / 1024 / 1024  # MB
-        memory_increase = final_memory - initial_memory
+        # Create new registry instance
+        registry2 = ModelRegistry(models_dir=temp_dir)
         
-        # Memory increase should be reasonable (less than 500MB for 10 images)
-        assert memory_increase < 500
+        # Check if model persisted
+        loaded_model = registry2.get_model("persistence-test")
+        assert loaded_model is not None
+        assert loaded_model.name == "Persistence Test"
+    
+    @pytest.mark.integration
+    def test_model_manager_workflow(self):
+        """Test complete model manager workflow."""
+        from models import create_model_manager
+        
+        manager = create_model_manager()
+        
+        # Test model discovery
+        stats = manager.get_stats()
+        assert "registry" in stats
+        assert stats["registry"]["total_models"] > 0
+        
+        # Test benchmark (without actual model loading)
+        with patch.object(manager.loader, 'load_model', return_value=Mock()):
+            benchmarks = manager.benchmark()
+            # Would return empty without real models
+            assert isinstance(benchmarks, dict)