conftest.py

import os
import tempfile
import shutil
from pathlib import Path
import pytest
import yaml
from unittest.mock import MagicMock


@pytest.fixture
def temp_dir():
    """Create a temporary directory for test files."""
    temp_path = tempfile.mkdtemp()
    yield Path(temp_path)
    shutil.rmtree(temp_path)


@pytest.fixture
def mock_config():
    """Create a mock configuration object."""
    config = MagicMock()
    config.MODEL = MagicMock()
    config.MODEL.TYPE = 'swin'
    config.MODEL.NAME = 'swin_base'
    config.MODEL.DROP_PATH_RATE = 0.1
    config.MODEL.NUM_CLASSES = 1000
    
    config.DATA = MagicMock()
    config.DATA.IMG_SIZE = 224
    config.DATA.BATCH_SIZE = 32
    config.DATA.DATA_PATH = '/path/to/data'
    
    config.TRAIN = MagicMock()
    config.TRAIN.EPOCHS = 100
    config.TRAIN.BASE_LR = 1e-4
    config.TRAIN.WEIGHT_DECAY = 0.05
    
    return config


@pytest.fixture
def sample_yaml_config(temp_dir):
    """Create a sample YAML configuration file."""
    config_data = {
        'MODEL': {
            'TYPE': 'vit',
            'NAME': 'vit_base',
            'NUM_CLASSES': 1000,
            'DROP_PATH_RATE': 0.1
        },
        'DATA': {
            'IMG_SIZE': 224,
            'BATCH_SIZE': 64,
            'DATA_PATH': '/data/imagenet'
        },
        'TRAIN': {
            'EPOCHS': 300,
            'BASE_LR': 5e-4,
            'WEIGHT_DECAY': 0.05
        }
    }
    
    config_path = temp_dir / 'test_config.yaml'
    with open(config_path, 'w') as f:
        yaml.dump(config_data, f)
    
    return config_path


@pytest.fixture
def mock_dataset():
    """Create a mock dataset object."""
    dataset = MagicMock()
    dataset.__len__ = MagicMock(return_value=1000)
    dataset.__getitem__ = MagicMock(return_value=(MagicMock(), 0))
    return dataset


@pytest.fixture
def mock_model():
    """Create a mock model object."""
    model = MagicMock()
    model.forward = MagicMock(return_value=MagicMock(shape=(32, 1000)))
    model.parameters = MagicMock(return_value=[MagicMock() for _ in range(10)])
    return model


@pytest.fixture
def sample_image_tensor():
    """Create a sample image tensor for testing."""
    try:
        import torch
        return torch.randn(1, 3, 224, 224)
    except ImportError:
        # Return a mock if torch is not available during testing
        return MagicMock(shape=(1, 3, 224, 224))


@pytest.fixture
def captured_output():
    """Capture stdout and stderr for testing print outputs."""
    import sys
    from io import StringIO
    
    old_stdout = sys.stdout
    old_stderr = sys.stderr
    sys.stdout = StringIO()
    sys.stderr = StringIO()
    
    yield sys.stdout, sys.stderr
    
    sys.stdout = old_stdout
    sys.stderr = old_stderr


@pytest.fixture(autouse=True)
def reset_modules():
    """Reset imported modules to ensure clean test state."""
    import sys
    modules_to_reset = [m for m in sys.modules.keys() if m.startswith(('models', 'data'))]
    for module in modules_to_reset:
        sys.modules.pop(module, None)
    yield