Create tests/test_models.py

tests/test_models.py

@@ -0,0 +1,349 @@
+"""
+Tests for the processing pipeline.
+"""
+
+import pytest
+import numpy as np
+import cv2
+from unittest.mock import Mock, patch, MagicMock
+from pathlib import Path
+
+from api.pipeline import (
+    ProcessingPipeline,
+    PipelineConfig,
+    PipelineResult,
+    ProcessingMode,
+    PipelineStage
+)
+
+
+class TestPipelineConfig:
+    """Test pipeline configuration."""
+    
+    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_custom_config(self):
+        """Test custom configuration."""
+        config = PipelineConfig(
+            mode=ProcessingMode.VIDEO,
+            quality_preset="ultra",
+            use_gpu=False,
+            batch_size=4
+        )
+        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."""
+    
+    @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_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_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_process_image_with_effects(self, mock_pipeline, sample_image):
+        """Test image processing with effects."""
+        mock_pipeline.config.apply_effects = ['bokeh', 'vignette']
+        
+        # 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
+    
+    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"))
+        
+        result = mock_pipeline.process_image(sample_image, None)
+        
+        assert result is not None
+        assert result.success == False
+        assert len(result.errors) > 0
+    
+    @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
+    
+    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
+        ))
+        
+        results = mock_pipeline.process_batch(images)
+        
+        assert len(results) == 3
+        assert all(r.success for r in results)
+    
+    def test_progress_callback(self, mock_pipeline, sample_image):
+        """Test progress callback functionality."""
+        progress_values = []
+        
+        def progress_callback(value, message):
+            progress_values.append(value)
+        
+        mock_pipeline.config.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
+        })
+        
+        result = mock_pipeline.process_image(sample_image, None)
+        
+        assert len(progress_values) > 0
+        assert 0.0 <= max(progress_values) <= 1.0
+    
+    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_memory_management(self, mock_pipeline):
+        """Test memory management and cleanup."""
+        initial_cache_size = len(mock_pipeline.cache)
+        
+        # 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)
+        
+        # Clear cache
+        mock_pipeline.clear_cache()
+        
+        assert len(mock_pipeline.cache) == 0
+    
+    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."""
+    
+    @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
+        )
+        
+        # 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()
+    
+    @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"
+        )
+        
+        try:
+            pipeline = ProcessingPipeline(config)
+        except Exception:
+            pytest.skip("Models not available for integration test")
+        
+        # Open video
+        cap = cv2.VideoCapture(sample_video)
+        processed_frames = []
+        
+        # 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()
+
+
+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
+    
+    @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
+        
+        # Process multiple images
+        for _ in range(10):
+            mock_pipeline.process_image(sample_image, None)
+        
+        final_memory = process.memory_info().rss / 1024 / 1024  # MB
+        memory_increase = final_memory - initial_memory
+        
+        # Memory increase should be reasonable (less than 500MB for 10 images)
+        assert memory_increase < 500