CLI implementation
This commit is contained in:
350
tests/unit/test_upscaler.py
Normal file
350
tests/unit/test_upscaler.py
Normal file
@@ -0,0 +1,350 @@
|
||||
"""
|
||||
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
|
||||
Reference in New Issue
Block a user