increase test coverage, pot. incorrect upsampler fix

Author: simon-donike

opensr_srgan/model/model_blocks/__init__.py

@@ -275,6 +275,9 @@ def _icnr_(weight: torch.Tensor, scale: int = 2) -> None:
 def make_upsampler(n_channels: int, scale: int, *, use_icnr: bool = False) -> nn.Sequential:
     """Create a pixel-shuffle upsampler matching the flexible generator implementation."""
 
+    if scale < 1 or (scale & (scale - 1)) != 0:
+        raise ValueError("scale must be a positive power of two (1, 2, 4, 8, ...).")
+
     stages: list[nn.Module] = []
     for _ in range(int(math.log2(scale))):
         conv = nn.Conv2d(n_channels, n_channels * 4, 3, padding=1)

tests/test_deployment_naming.py

@@ -0,0 +1,26 @@
+from __future__ import annotations
+
+from pathlib import Path
+
+from deployment.srgan_hpc.naming import (
+    fused_output_name,
+    patch_dir,
+    patch_output_name,
+    product_output_name,
+    resolve_run_dir,
+)
+
+
+def test_output_name_helpers_format_expected_names() -> None:
+    assert patch_output_name(45.1234567, 9.9876543) == "output_SR_image_45.123457_9.987654.tif"
+    assert product_output_name("rgbnir") == "rgbnir_sr.tif"
+    assert fused_output_name() == "fused_sr.tif"
+
+
+def test_run_and_patch_directory_helpers_join_paths() -> None:
+    output_root = Path("/tmp/srgan-runs")
+    run_dir = resolve_run_dir(output_root, "run_001")
+    assert run_dir == output_root / "run_001"
+
+    patch_path = patch_dir(run_dir, "patch_0001")
+    assert patch_path == output_root / "run_001" / "patches" / "patch_0001"

tests/test_model_blocks/test_blocks.py

@@ -0,0 +1,68 @@
+from __future__ import annotations
+
+import pytest
+
+torch = pytest.importorskip("torch")
+
+from opensr_srgan.model.model_blocks import (
+    ConvolutionalBlock,
+    DenseBlock5,
+    LKA,
+    LKAResBlock,
+    RCAB,
+    RRDB,
+    ResidualBlock,
+    ResidualBlockNoBN,
+    SubPixelConvolutionalBlock,
+    make_upsampler,
+)
+from opensr_srgan.model.model_blocks import _icnr_
+
+
+def test_convolutional_block_rejects_unknown_activation() -> None:
+    with pytest.raises(AssertionError, match="activation must be one of"):
+        ConvolutionalBlock(16, 16, 3, activation="relu")
+
+
+def test_convolutional_block_tanh_path_keeps_spatial_shape() -> None:
+    block = ConvolutionalBlock(16, 16, 3, batch_norm=True, activation="tanh")
+    x = torch.randn(2, 16, 8, 8)
+    y = block(x)
+    assert y.shape == x.shape
+
+
+def test_subpixel_block_upsamples_by_scaling_factor() -> None:
+    block = SubPixelConvolutionalBlock(n_channels=16, scaling_factor=2)
+    x = torch.randn(1, 16, 8, 8)
+    y = block(x)
+    assert y.shape == (1, 16, 16, 16)
+
+
+def test_residual_and_attention_blocks_preserve_shape() -> None:
+    x = torch.randn(1, 16, 8, 8)
+
+    assert ResidualBlock(n_channels=16)(x).shape == x.shape
+    assert ResidualBlockNoBN(n_channels=16)(x).shape == x.shape
+    assert RCAB(n_channels=16)(x).shape == x.shape
+    assert DenseBlock5(n_features=16, growth_channels=8)(x).shape == x.shape
+    assert RRDB(n_features=16, growth_channels=8)(x).shape == x.shape
+    assert LKA(n_channels=16)(x).shape == x.shape
+    assert LKAResBlock(n_channels=16)(x).shape == x.shape
+
+
+def test_icnr_requires_divisible_output_channels() -> None:
+    weight = torch.empty(10, 4, 3, 3)
+    with pytest.raises(ValueError, match=r"divisible by scale\*\*2"):
+        _icnr_(weight, scale=2)
+
+
+def test_make_upsampler_scale_4_with_icnr_produces_expected_shape() -> None:
+    upsampler = make_upsampler(16, scale=4, use_icnr=True)
+    x = torch.randn(1, 16, 8, 8)
+    y = upsampler(x)
+    assert y.shape == (1, 16, 32, 32)
+
+
+def test_make_upsampler_rejects_non_power_of_two_scale() -> None:
+    with pytest.raises(ValueError, match="power of two"):
+        make_upsampler(16, scale=3)