tests/conftest.py

"""
Pytest configuration and fixtures for BackgroundFX Pro tests.
"""

import pytest
import numpy as np
import torch
import cv2
import tempfile
import shutil
from pathlib import Path
from unittest.mock import Mock, MagicMock
import os
import sys

# Add parent directory to path for imports
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))


# ============================================================================
# Configuration
# ============================================================================

@pytest.fixture(scope="session")
def test_config():
    """Test configuration."""
    return {
        'device': 'cpu',  # Use CPU for testing
        'test_data_dir': Path(__file__).parent / 'data',
        'temp_dir': tempfile.mkdtemp(prefix='bgfx_test_'),
        'max_test_duration': 30,  # seconds
        'use_gpu': torch.cuda.is_available()
    }


@pytest.fixture(scope="session", autouse=True)
def cleanup(test_config):
    """Cleanup after all tests."""
    yield
    # Clean up temp directory
    if os.path.exists(test_config['temp_dir']):
        shutil.rmtree(test_config['temp_dir'])


# ============================================================================
# Image and Video Fixtures
# ============================================================================

@pytest.fixture
def sample_image():
    """Create a sample image for testing."""
    # Create 512x512 RGB image with a person-like shape
    image = np.zeros((512, 512, 3), dtype=np.uint8)
    
    # Add background
    image[:, :] = [100, 150, 200]  # Blue background
    
    # Add person-like shape (simple rectangle for testing)
    cv2.rectangle(image, (150, 100), (350, 450), (50, 100, 50), -1)
    
    # Add some texture
    noise = np.random.randint(0, 20, (512, 512, 3), dtype=np.uint8)
    image = cv2.add(image, noise)
    
    return image


@pytest.fixture
def sample_mask():
    """Create a sample mask for testing."""
    mask = np.zeros((512, 512), dtype=np.uint8)
    # Create person mask
    cv2.rectangle(mask, (150, 100), (350, 450), 255, -1)
    # Add some edge refinement
    mask = cv2.GaussianBlur(mask, (5, 5), 2)
    return mask


@pytest.fixture
def sample_background():
    """Create a sample background image."""
    background = np.zeros((512, 512, 3), dtype=np.uint8)
    # Create gradient background
    for i in range(512):
        background[i, :] = [
            int(255 * (i / 512)),  # Red gradient
            100,  # Fixed green
            int(255 * (1 - i / 512))  # Blue inverse gradient
        ]
    return background


@pytest.fixture
def sample_video(test_config):
    """Create a sample video file for testing."""
    video_path = Path(test_config['temp_dir']) / 'test_video.mp4'
    
    # Create video writer
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    out = cv2.VideoWriter(str(video_path), fourcc, 30.0, (512, 512))
    
    # Write 30 frames (1 second at 30fps)
    for i in range(30):
        frame = np.zeros((512, 512, 3), dtype=np.uint8)
        # Animate a moving rectangle
        x = 100 + i * 5
        cv2.rectangle(frame, (x, 200), (x + 100, 400), (0, 255, 0), -1)
        out.write(frame)
    
    out.release()
    return str(video_path)


# ============================================================================
# Model Fixtures
# ============================================================================

@pytest.fixture
def mock_model():
    """Create a mock ML model for testing."""
    model = MagicMock()
    model.eval = MagicMock(return_value=None)
    model.to = MagicMock(return_value=model)
    
    # Mock forward pass
    def forward(x):
        batch_size = x.shape[0] if hasattr(x, 'shape') else 1
        return torch.randn(batch_size, 1, 512, 512)
    
    model.__call__ = MagicMock(side_effect=forward)
    model.forward = MagicMock(side_effect=forward)
    
    return model


@pytest.fixture
def mock_sam2_predictor():
    """Create a mock SAM2 predictor."""
    predictor = MagicMock()
    
    def predict(image):
        h, w = image.shape[:2] if len(image.shape) > 2 else (512, 512)
        return np.random.randint(0, 2, (h, w), dtype=np.uint8) * 255
    
    predictor.predict = MagicMock(side_effect=predict)
    predictor.set_image = MagicMock(return_value=None)
    
    return predictor


@pytest.fixture
def mock_matanyone_model():
    """Create a mock MatAnyone model."""
    model = MagicMock()
    
    def refine(image, mask):
        return cv2.GaussianBlur(mask, (5, 5), 2)
    
    model.refine = MagicMock(side_effect=refine)
    
    return model


# ============================================================================
# Pipeline and Processing Fixtures
# ============================================================================

@pytest.fixture
def pipeline_config():
    """Create pipeline configuration for testing."""
    from api.pipeline import PipelineConfig
    
    return PipelineConfig(
        use_gpu=False,  # CPU for testing
        quality_preset='medium',
        enable_cache=False,  # Disable cache for testing
        batch_size=1,
        max_workers=2
    )


@pytest.fixture
def mock_pipeline(pipeline_config):
    """Create a mock processing pipeline."""
    from api.pipeline import ProcessingPipeline
    
    # Mock the pipeline to avoid loading real models
    with pytest.MonkeyPatch().context() as m:
        m.setattr('api.pipeline.ModelFactory.load_model', 
                  lambda self, *args, **kwargs: Mock())
        pipeline = ProcessingPipeline(pipeline_config)
    
    return pipeline


# ============================================================================
# API and Server Fixtures
# ============================================================================

@pytest.fixture
def api_client():
    """Create a test client for the API."""
    from fastapi.testclient import TestClient
    from api.api_server import app
    
    return TestClient(app)


@pytest.fixture
def mock_job_manager():
    """Create a mock job manager."""
    manager = MagicMock()
    manager.create_job = MagicMock(return_value='test-job-123')
    manager.get_job = MagicMock(return_value={'status': 'processing'})
    manager.update_job = MagicMock(return_value=None)
    
    return manager


# ============================================================================
# File System Fixtures
# ============================================================================

@pytest.fixture
def temp_dir(test_config):
    """Create a temporary directory for test files."""
    temp_path = Path(test_config['temp_dir']) / 'test_run'
    temp_path.mkdir(parents=True, exist_ok=True)
    yield temp_path
    # Cleanup
    if temp_path.exists():
        shutil.rmtree(temp_path)


@pytest.fixture
def sample_files(temp_dir, sample_image):
    """Create sample files in temp directory."""
    files = {}
    
    # Save sample image
    image_path = temp_dir / 'sample.jpg'
    cv2.imwrite(str(image_path), sample_image)
    files['image'] = image_path
    
    # Create multiple images for batch testing
    for i in range(3):
        path = temp_dir / f'image_{i}.jpg'
        cv2.imwrite(str(path), sample_image)
        files[f'image_{i}'] = path
    
    return files


# ============================================================================
# Model Registry Fixtures
# ============================================================================

@pytest.fixture
def mock_registry():
    """Create a mock model registry."""
    from models.registry import ModelRegistry, ModelInfo, ModelTask, ModelFramework
    
    registry = ModelRegistry(models_dir=Path(tempfile.mkdtemp()))
    
    # Add test model
    test_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_model.pth',
        file_size=1000000
    )
    
    registry.register_model(test_model)
    
    return registry


# ============================================================================
# WebSocket Fixtures
# ============================================================================

@pytest.fixture
def mock_websocket():
    """Create a mock WebSocket connection."""
    ws = MagicMock()
    ws.accept = MagicMock(return_value=None)
    ws.send_json = MagicMock(return_value=None)
    ws.receive_text = MagicMock(return_value='{"type": "ping", "data": {}}')
    
    return ws


# ============================================================================
# Utility Fixtures
# ============================================================================

@pytest.fixture
def mock_progress_callback():
    """Create a mock progress callback."""
    callback = MagicMock()
    return callback


@pytest.fixture
def device():
    """Get device for testing."""
    return 'cuda' if torch.cuda.is_available() else 'cpu'


@pytest.fixture
def performance_timer():
    """Timer for performance testing."""
    import time
    
    class Timer:
        def __init__(self):
            self.start_time = None
            self.elapsed = 0
        
        def __enter__(self):
            self.start_time = time.time()
            return self
        
        def __exit__(self, *args):
            self.elapsed = time.time() - self.start_time
    
    return Timer


# ============================================================================
# Markers
# ============================================================================

def pytest_configure(config):
    """Register custom markers."""
    config.addinivalue_line(
        "markers", "slow: marks tests as slow (deselect with '-m \"not slow\"')"
    )
    config.addinivalue_line(
        "markers", "gpu: marks tests that require GPU"
    )
    config.addinivalue_line(
        "markers", "integration: marks integration tests"
    )
    config.addinivalue_line(
        "markers", "unit: marks unit tests"
    )