351 lines
13 KiB
Python
351 lines
13 KiB
Python
"""
|
|
Tests for upscaling module.
|
|
"""
|
|
|
|
import pytest
|
|
import tempfile
|
|
from pathlib import Path
|
|
from unittest.mock import patch, MagicMock, mock_open
|
|
from src.upscaling.upscaler import (
|
|
UpscaleConfig,
|
|
UpscaleResult,
|
|
UpscalerType,
|
|
UpscaleFactor,
|
|
FFmpegSRUpscaler,
|
|
RealESRGANUpscaler,
|
|
UpscaleManager,
|
|
BaseUpscaler
|
|
)
|
|
|
|
|
|
class TestUpscaleConfig:
|
|
"""Test UpscaleConfig dataclass."""
|
|
|
|
def test_default_config(self):
|
|
"""Test default configuration values."""
|
|
config = UpscaleConfig()
|
|
assert config.upscaler_type == UpscalerType.FFMPEG_SR
|
|
assert config.factor == UpscaleFactor.X2
|
|
assert config.denoise_strength == 0.5
|
|
assert config.tile_size == 0
|
|
assert config.half_precision is True
|
|
assert config.device == "cuda"
|
|
|
|
def test_custom_config(self):
|
|
"""Test custom configuration."""
|
|
config = UpscaleConfig(
|
|
upscaler_type=UpscalerType.REAL_ESRGAN,
|
|
factor=UpscaleFactor.X4,
|
|
denoise_strength=0.8,
|
|
device="cpu"
|
|
)
|
|
assert config.upscaler_type == UpscalerType.REAL_ESRGAN
|
|
assert config.factor == UpscaleFactor.X4
|
|
assert config.denoise_strength == 0.8
|
|
assert config.device == "cpu"
|
|
|
|
def test_config_to_dict(self):
|
|
"""Test configuration serialization."""
|
|
config = UpscaleConfig(
|
|
factor=UpscaleFactor.X4,
|
|
model_path=Path("/models/model.pth")
|
|
)
|
|
d = config.to_dict()
|
|
assert d["factor"] == 4
|
|
assert d["upscaler_type"] == "ffmpeg_sr"
|
|
# Path conversion is platform-specific
|
|
assert "model.pth" in d["model_path"]
|
|
|
|
|
|
class TestUpscaleResult:
|
|
"""Test UpscaleResult dataclass."""
|
|
|
|
def test_success_result(self):
|
|
"""Test successful result."""
|
|
result = UpscaleResult(
|
|
success=True,
|
|
output_path=Path("output.mp4"),
|
|
input_resolution=(1920, 1080),
|
|
output_resolution=(3840, 2160),
|
|
processing_time_s=45.5
|
|
)
|
|
assert result.success is True
|
|
assert result.output_path == Path("output.mp4")
|
|
assert result.input_resolution == (1920, 1080)
|
|
assert result.output_resolution == (3840, 2160)
|
|
assert result.processing_time_s == 45.5
|
|
|
|
def test_failure_result(self):
|
|
"""Test failure result."""
|
|
result = UpscaleResult(
|
|
success=False,
|
|
error_message="FFmpeg not found"
|
|
)
|
|
assert result.success is False
|
|
assert result.error_message == "FFmpeg not found"
|
|
|
|
|
|
class TestFFmpegSRUpscaler:
|
|
"""Test FFmpeg-based upscaler."""
|
|
|
|
@pytest.fixture
|
|
def mock_subprocess(self):
|
|
"""Fixture to mock subprocess.run."""
|
|
with patch('subprocess.run') as mock_run:
|
|
mock_run.return_value = MagicMock(returncode=0, stdout="", stderr="")
|
|
yield mock_run
|
|
|
|
@pytest.fixture
|
|
def temp_dir(self):
|
|
"""Fixture for temporary directory."""
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
yield Path(tmpdir)
|
|
|
|
def test_is_available_success(self, mock_subprocess):
|
|
"""Test availability check when ffmpeg is present."""
|
|
config = UpscaleConfig()
|
|
upscaler = FFmpegSRUpscaler(config)
|
|
assert upscaler.is_available() is True
|
|
|
|
def test_is_available_failure(self):
|
|
"""Test availability check when ffmpeg is missing."""
|
|
with patch('subprocess.run', side_effect=FileNotFoundError()):
|
|
config = UpscaleConfig()
|
|
upscaler = FFmpegSRUpscaler(config)
|
|
assert upscaler.is_available() is False
|
|
|
|
def test_upscale_missing_input(self, mock_subprocess):
|
|
"""Test upscaling with missing input file."""
|
|
config = UpscaleConfig()
|
|
upscaler = FFmpegSRUpscaler(config)
|
|
|
|
result = upscaler.upscale(Path("nonexistent.mp4"), Path("output.mp4"))
|
|
|
|
assert result.success is False
|
|
assert "not found" in result.error_message
|
|
|
|
def test_upscale_success(self, mock_subprocess, temp_dir):
|
|
"""Test successful upscaling."""
|
|
# Create dummy input file
|
|
input_file = temp_dir / "input.mp4"
|
|
input_file.write_bytes(b"dummy video")
|
|
output_file = temp_dir / "output.mp4"
|
|
|
|
# Mock ffprobe response
|
|
ffprobe_response = MagicMock(
|
|
returncode=0,
|
|
stdout='{"streams": [{"width": 1920, "height": 1080, "r_frame_rate": "24/1", "duration": "10.0"}]}'
|
|
)
|
|
|
|
with patch('subprocess.run') as mock_run:
|
|
# First call is ffprobe, second is ffmpeg
|
|
mock_run.side_effect = [ffprobe_response, MagicMock(returncode=0)]
|
|
|
|
config = UpscaleConfig(factor=UpscaleFactor.X2)
|
|
upscaler = FFmpegSRUpscaler(config)
|
|
result = upscaler.upscale(input_file, output_file)
|
|
|
|
assert result.success is True
|
|
assert result.input_resolution == (1920, 1080)
|
|
assert result.output_resolution == (3840, 2160)
|
|
|
|
def test_upscale_ffmpeg_error(self, mock_subprocess, temp_dir):
|
|
"""Test handling of ffmpeg error."""
|
|
input_file = temp_dir / "input.mp4"
|
|
input_file.write_bytes(b"dummy video")
|
|
output_file = temp_dir / "output.mp4"
|
|
|
|
ffprobe_response = MagicMock(
|
|
returncode=0,
|
|
stdout='{"streams": [{"width": 1920, "height": 1080, "r_frame_rate": "24/1", "duration": "10.0"}]}'
|
|
)
|
|
|
|
ffmpeg_error = MagicMock(returncode=1, stderr="Invalid data")
|
|
|
|
with patch('subprocess.run') as mock_run:
|
|
mock_run.side_effect = [ffprobe_response, ffmpeg_error]
|
|
|
|
config = UpscaleConfig()
|
|
upscaler = FFmpegSRUpscaler(config)
|
|
result = upscaler.upscale(input_file, output_file)
|
|
|
|
assert result.success is False
|
|
assert "FFmpeg error" in result.error_message
|
|
|
|
|
|
class TestRealESRGANUpscaler:
|
|
"""Test Real-ESRGAN upscaler."""
|
|
|
|
def test_is_available_without_realesrgan(self):
|
|
"""Test availability when Real-ESRGAN is not installed."""
|
|
with patch.dict('sys.modules', {'realesrgan': None}):
|
|
config = UpscaleConfig()
|
|
upscaler = RealESRGANUpscaler(config)
|
|
assert upscaler.is_available() is False
|
|
|
|
def test_upscale_without_realesrgan(self):
|
|
"""Test upscaling when Real-ESRGAN is not installed."""
|
|
# Mock the import to fail
|
|
import sys
|
|
original_modules = sys.modules.copy()
|
|
|
|
try:
|
|
# Remove realesrgan and torch from modules to simulate not installed
|
|
for mod in list(sys.modules.keys()):
|
|
if 'realesrgan' in mod or 'torch' in mod:
|
|
del sys.modules[mod]
|
|
|
|
# Mock import to raise ImportError
|
|
def mock_import(name, *args, **kwargs):
|
|
if name == 'realesrgan' or name.startswith('realesrgan.'):
|
|
raise ImportError("No module named 'realesrgan'")
|
|
if name == 'torch' or name.startswith('torch.'):
|
|
raise ImportError("No module named 'torch'")
|
|
return original_modules.get(name, __import__(name, *args, **kwargs))
|
|
|
|
with patch('builtins.__import__', side_effect=mock_import):
|
|
config = UpscaleConfig()
|
|
upscaler = RealESRGANUpscaler(config)
|
|
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
input_file = Path(tmpdir) / "input.mp4"
|
|
input_file.write_bytes(b"dummy")
|
|
output_file = Path(tmpdir) / "output.mp4"
|
|
|
|
result = upscaler.upscale(input_file, output_file)
|
|
|
|
assert result.success is False
|
|
assert "not installed" in result.error_message
|
|
finally:
|
|
# Restore original modules
|
|
sys.modules.clear()
|
|
sys.modules.update(original_modules)
|
|
|
|
|
|
class TestUpscaleManager:
|
|
"""Test UpscaleManager."""
|
|
|
|
def test_init_with_default_config(self):
|
|
"""Test initialization with default config."""
|
|
manager = UpscaleManager()
|
|
assert manager.config.upscaler_type == UpscalerType.FFMPEG_SR
|
|
|
|
def test_init_with_custom_config(self):
|
|
"""Test initialization with custom config."""
|
|
config = UpscaleConfig(upscaler_type=UpscalerType.REAL_ESRGAN)
|
|
manager = UpscaleManager(config)
|
|
assert manager.config.upscaler_type == UpscalerType.REAL_ESRGAN
|
|
|
|
def test_get_upscaler_caching(self):
|
|
"""Test that upscalers are cached."""
|
|
with patch.object(FFmpegSRUpscaler, 'is_available', return_value=True):
|
|
config = UpscaleConfig()
|
|
manager = UpscaleManager(config)
|
|
|
|
upscaler1 = manager.get_upscaler(UpscalerType.FFMPEG_SR)
|
|
upscaler2 = manager.get_upscaler(UpscalerType.FFMPEG_SR)
|
|
|
|
assert upscaler1 is upscaler2 # Same instance
|
|
|
|
def test_get_upscaler_unavailable(self):
|
|
"""Test getting unavailable upscaler."""
|
|
with patch.object(FFmpegSRUpscaler, 'is_available', return_value=False):
|
|
config = UpscaleConfig()
|
|
manager = UpscaleManager(config)
|
|
|
|
with pytest.raises(RuntimeError) as exc_info:
|
|
manager.get_upscaler(UpscalerType.FFMPEG_SR)
|
|
|
|
assert "not available" in str(exc_info.value)
|
|
|
|
def test_upscale_with_manager(self):
|
|
"""Test upscale through manager."""
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
input_file = Path(tmpdir) / "input.mp4"
|
|
input_file.write_bytes(b"dummy")
|
|
output_file = Path(tmpdir) / "output.mp4"
|
|
|
|
mock_result = UpscaleResult(
|
|
success=True,
|
|
output_path=output_file,
|
|
input_resolution=(1920, 1080),
|
|
output_resolution=(3840, 2160)
|
|
)
|
|
|
|
with patch.object(FFmpegSRUpscaler, 'is_available', return_value=True):
|
|
with patch.object(FFmpegSRUpscaler, 'upscale', return_value=mock_result):
|
|
config = UpscaleConfig()
|
|
manager = UpscaleManager(config)
|
|
result = manager.upscale(input_file, output_file)
|
|
|
|
assert result.success is True
|
|
assert result.output_resolution == (3840, 2160)
|
|
|
|
def test_get_available_upscalers(self):
|
|
"""Test getting list of available upscalers."""
|
|
with patch.object(FFmpegSRUpscaler, 'is_available', return_value=True):
|
|
with patch.object(RealESRGANUpscaler, 'is_available', return_value=False):
|
|
config = UpscaleConfig()
|
|
manager = UpscaleManager(config)
|
|
available = manager.get_available_upscalers()
|
|
|
|
assert UpscalerType.FFMPEG_SR in available
|
|
assert UpscalerType.REAL_ESRGAN not in available
|
|
|
|
def test_upscale_with_fallback_success(self):
|
|
"""Test fallback upscaling with first option succeeding."""
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
input_file = Path(tmpdir) / "input.mp4"
|
|
input_file.write_bytes(b"dummy")
|
|
output_file = Path(tmpdir) / "output.mp4"
|
|
|
|
mock_result = UpscaleResult(
|
|
success=True,
|
|
output_path=output_file,
|
|
input_resolution=(1920, 1080),
|
|
output_resolution=(3840, 2160)
|
|
)
|
|
|
|
with patch.object(FFmpegSRUpscaler, 'is_available', return_value=True):
|
|
with patch.object(FFmpegSRUpscaler, 'upscale', return_value=mock_result):
|
|
config = UpscaleConfig()
|
|
manager = UpscaleManager(config)
|
|
result = manager.upscale_with_fallback(input_file, output_file)
|
|
|
|
assert result.success is True
|
|
|
|
def test_upscale_with_fallback_all_fail(self):
|
|
"""Test fallback upscaling when all options fail."""
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
input_file = Path(tmpdir) / "input.mp4"
|
|
input_file.write_bytes(b"dummy")
|
|
output_file = Path(tmpdir) / "output.mp4"
|
|
|
|
mock_result = UpscaleResult(
|
|
success=False,
|
|
error_message="Processing failed"
|
|
)
|
|
|
|
with patch.object(FFmpegSRUpscaler, 'is_available', return_value=True):
|
|
with patch.object(FFmpegSRUpscaler, 'upscale', return_value=mock_result):
|
|
config = UpscaleConfig()
|
|
manager = UpscaleManager(config)
|
|
result = manager.upscale_with_fallback(input_file, output_file)
|
|
|
|
assert result.success is False
|
|
assert "All upscalers failed" in result.error_message
|
|
|
|
|
|
class TestUpscaleEnums:
|
|
"""Test upscaler enums."""
|
|
|
|
def test_upscaler_type_values(self):
|
|
"""Test upscaler type enum values."""
|
|
assert UpscalerType.REAL_ESRGAN.value == "real_esrgan"
|
|
assert UpscalerType.FFMPEG_SR.value == "ffmpeg_sr"
|
|
|
|
def test_upscale_factor_values(self):
|
|
"""Test upscale factor enum values."""
|
|
assert UpscaleFactor.X2.value == 2
|
|
assert UpscaleFactor.X4.value == 4
|