Add tests

Author: Sllambias

.github/workflows/pipeline_tests.yml

@@ -0,0 +1,34 @@
+name: Pipeline tests
+
+on:
+  push:
+    branches: [main]
+  pull_request:
+    branches: [main]
+
+jobs:
+  pipeline-tests:
+    runs-on: ubuntu-latest
+    timeout-minutes: 20
+
+    steps:
+      - uses: actions/checkout@v4
+
+      - uses: actions/setup-python@v5
+        with:
+          python-version: "3.11"
+          cache: pip
+
+      - name: Install CPU-only torch then package
+        run: |
+          pip install torch torchvision --index-url https://download.pytorch.org/whl/cpu
+          pip install -e ".[test]"
+
+      - name: Run pipeline tests
+        run: |
+          pytest tests/test_pretrain.py \
+                 tests/test_finetune_seg.py \
+                 tests/test_train_reg.py \
+                 tests/test_test_cls.py \
+                 tests/test_linear_probe.py \
+                 -v --timeout=300

pyproject.toml

@@ -58,7 +58,8 @@ extras = [
 ]
 test = [
     "ruff>=0.14.8",
-    "pytest>=9.0.1"
+    "pytest>=9.0.1",
+    "pytest-timeout>=2.3.0",
 ]
 docs = [
     "mkdocs-shadcn",

tests/conftest.py

@@ -0,0 +1,103 @@
+import pickle
+import pytest
+import torch
+from lightning import Trainer
+
+
+@pytest.fixture
+def pretrain_files(tmp_path):
+    """Three .pt files of shape [1, 32, 32, 32] for pretraining (raw image, no label).
+    32^3 ensures the UNet bottleneck (4 max-pool stages) stays at 2x2x2, avoiding
+    single-element BatchNorm errors with batch_size=1.
+    """
+    files = []
+    for i in range(3):
+        path = tmp_path / f"pre_{i:03d}.pt"
+        torch.save(torch.randn(1, 32, 32, 32), path)
+        files.append(str(path))
+    return {"train": files[:2], "val": [files[2]]}
+
+
+@pytest.fixture
+def seg_files(tmp_path):
+    """Three .pt + .pkl file pairs for segmentation. Shape [2, 32, 32, 32] = [image, label].
+    32^3 ensures the UNet bottleneck (4 max-pool stages) stays at 2x2x2.
+    """
+    files = []
+    for i in range(3):
+        pt = tmp_path / f"seg_{i:03d}.pt"
+        pkl = tmp_path / f"seg_{i:03d}.pkl"
+        data = torch.zeros(2, 32, 32, 32)
+        data[0] = torch.randn(32, 32, 32)
+        data[1] = torch.randint(0, 2, (32, 32, 32)).float()
+        torch.save(data, pt)
+        with open(pkl, "wb") as f:
+            pickle.dump({"foreground_locations": []}, f)
+        files.append(str(pt))
+    return {"train": files[:2], "val": [files[2]]}
+
+
+@pytest.fixture
+def clsreg_files(tmp_path):
+    """Three .pt files containing (image[1,32,32,32], label_scalar) tuples.
+    32^3 prevents single-element BatchNorm errors in the 4-stage UNet encoder.
+    Labels are 0-dim int tensors; ClassificationModule.on_before_batch_transfer
+    squeezes and converts to long before the training step.
+    """
+    files = []
+    for i in range(3):
+        path = tmp_path / f"cls_{i:03d}.pt"
+        torch.save((torch.randn(1, 32, 32, 32), torch.tensor(i % 2)), path)
+        files.append(str(path))
+    return {"train": files[:2], "val": [files[2]], "test": [files[2]]}
+
+
+@pytest.fixture
+def reg_files(tmp_path):
+    """Three .pt files containing (image[1,32,32,32], label[1]) tuples.
+    Labels are 1D float tensors so they collate to [B, 1], matching the
+    unet_clsreg_tiny output shape [B, 1] expected by MeanSquaredError.
+    """
+    files = []
+    for i in range(3):
+        path = tmp_path / f"reg_{i:03d}.pt"
+        torch.save((torch.randn(1, 32, 32, 32), torch.tensor([float(i % 2)])), path)
+        files.append(str(path))
+    return {"train": files[:2], "val": [files[2]], "test": [files[2]]}
+
+
+@pytest.fixture
+def cls_probe_files(tmp_path):
+    """Five .pt files for classification / linear-probe tests. 0-dim integer labels.
+    2 train + 2 val gives full batches when batch_size=2, avoiding the squeeze()-to-scalar
+    edge case in ClassificationModule.on_before_batch_transfer with batch_size=1.
+    2 test files (labels 1, 0) ensure both classes are present for AUROC computation.
+    """
+    labels = [0, 1, 0, 1, 0, 1]
+    files = []
+    for i, lbl in enumerate(labels):
+        path = tmp_path / f"clsp_{i:03d}.pt"
+        torch.save((torch.randn(1, 32, 32, 32), torch.tensor(lbl)), path)
+        files.append(str(path))
+    return {"train": files[:2], "val": files[2:4], "test": files[4:6]}
+
+
+@pytest.fixture
+def make_trainer(tmp_path):
+    """Factory fixture that builds a minimal CPU Trainer for smoke tests."""
+
+    def _make(**kwargs):
+        defaults = dict(
+            accelerator="cpu",
+            max_epochs=1,
+            limit_train_batches=2,
+            limit_val_batches=2,
+            logger=False,
+            enable_checkpointing=False,
+            enable_progress_bar=False,
+            num_sanity_val_steps=0,
+        )
+        defaults.update(kwargs)
+        return Trainer(default_root_dir=str(tmp_path), **defaults)
+
+    return _make

tests/test_finetune_seg.py

@@ -0,0 +1,33 @@
+"""Integration test for pipeline/run/finetune_seg.py components.
+
+Uses SegmentationModule + SegDataModule + unet_tiny on synthetic 8x8x8 seg volumes.
+Only tests trainer.fit() — trainer.test() is excluded because SegTestDataset._get_src_label()
+loads from ASPARAGUS_RAW_LABELS which is unavailable in CI.
+"""
+from asparagus.modules.data_modules.training import SegDataModule
+from asparagus.modules.lightning_modules import SegmentationModule
+from asparagus.modules.networks.unet import unet_tiny
+
+
+def test_finetune_seg_fit(seg_files, make_trainer):
+    """SegmentationModule fits from scratch (weights=None) on synthetic seg data."""
+    model = unet_tiny(input_channels=1, output_channels=2, dimensions="3D")
+
+    data_module = SegDataModule(
+        batch_size=1,
+        num_workers=1,
+        train_split=seg_files["train"],
+        val_split=seg_files["val"],
+        train_transforms=None,
+        val_transforms=None,
+    )
+
+    module = SegmentationModule(
+        model=model,
+        learning_rate=1e-3,
+        warmup_epochs=0,
+        weights=None,
+        inference_patch_size=[32, 32, 32],
+    )
+
+    make_trainer().fit(module, datamodule=data_module)

tests/test_linear_probe.py

@@ -0,0 +1,53 @@
+"""Integration test for pipeline/run/linear_probe.py components.
+
+Mirrors linear_probe.py's three-phase flow: validate → fit → test.
+Uses LinearProbeModule + ClsRegDataModule + ResidualEncoderUNetCLSREG (tiny).
+
+LinearProbeModule calls model._encode() internally, which is defined on
+ResidualEncoderUNetCLSREG but not on UNetCLSREG.
+
+batch_size=2 is required: squeeze(-1) in on_before_batch_transfer collapses
+a [B] label tensor to 0-dim when B=1, causing CrossEntropyLoss to fail.
+limit_test_batches=2 ensures both test files (labels 0 and 1) are processed
+so MulticlassAUROC has both classes present.
+"""
+from asparagus.modules.data_modules.training import ClsRegDataModule
+from asparagus.modules.lightning_modules import LinearProbeModule
+from asparagus.modules.networks.resenc_unet import ResidualEncoderUNetCLSREG
+
+
+def test_linear_probe_validate_fit_test(cls_probe_files, tmp_path, make_trainer):
+    """LinearProbeModule runs all three phases: validate → fit → test."""
+    model = ResidualEncoderUNetCLSREG(
+        input_channels=1,
+        output_channels=2,
+        dimensions="3D",
+        features_per_stage=(4, 8),
+        stride=2,
+        kernel_size=3,
+        n_blocks_per_stage=(1, 1),
+    )
+
+    data_module = ClsRegDataModule(
+        batch_size=2,
+        num_workers=2,  # val_dataloader uses num_workers//2; needs >=2
+        train_split=cls_probe_files["train"],
+        val_split=cls_probe_files["val"],
+        test_samples=cls_probe_files["test"],
+        use_random_datasampler=False,
+    )
+
+    module = LinearProbeModule(
+        model=model,
+        learning_rates=[0.1, 0.01],
+        num_classes=2,
+        dimensions="3D",
+        test_output_path=str(tmp_path / "probe_preds.json"),
+        weights=None,
+    )
+
+    trainer = make_trainer(limit_test_batches=2)
+    data_module.setup("fit")
+    trainer.validate(module, datamodule=data_module)
+    trainer.fit(module, datamodule=data_module)
+    trainer.test(module, datamodule=data_module)

tests/test_pretrain.py

@@ -0,0 +1,39 @@
+"""Integration test for pipeline/run/pretrain.py components.
+
+Uses SelfSupervisedModule + PretrainDataModule + unet_tiny on synthetic 8x8x8 volumes.
+Torch_CopyImageToLabel adds batch["label"] so the SSL reconstruction loss can run.
+"""
+from gardening_tools.modules.transforms.copy_image_to_label import Torch_CopyImageToLabel
+from torchvision import transforms
+
+from asparagus.modules.data_modules.pretraining import PretrainDataModule
+from asparagus.modules.lightning_modules import SelfSupervisedModule
+from asparagus.modules.networks.unet import unet_tiny
+
+
+def test_pretrain_fit(pretrain_files, make_trainer):
+    """SelfSupervisedModule fits on synthetic pretrain data with reconstruction loss."""
+    model = unet_tiny(input_channels=1, output_channels=1, dimensions="3D")
+
+    # CopyImageToLabel saves label = image before any GPU augmentation,
+    # which is all the SSL reconstruction loss requires.
+    copy_transform = transforms.Compose([Torch_CopyImageToLabel(copy=True)])
+
+    data_module = PretrainDataModule(
+        batch_size=1,
+        num_workers=1,
+        train_split=pretrain_files["train"],
+        val_split=pretrain_files["val"],
+        train_transforms=copy_transform,
+        val_transforms=copy_transform,
+    )
+
+    module = SelfSupervisedModule(
+        model=model,
+        learning_rate=1e-3,
+        warmup_epochs=0,
+        train_transforms=None,
+        val_transforms=None,
+    )
+
+    make_trainer().fit(module, datamodule=data_module)

tests/test_test_cls.py

@@ -0,0 +1,57 @@
+"""Integration test for pipeline/run/test_cls.py components.
+
+Mirrors test_cls.py's two-stage flow:
+  1. Train briefly, save a checkpoint.
+  2. Load that checkpoint, run test-time inference on new data.
+Uses ClassificationModule + ClsRegDataModule + unet_clsreg_tiny.
+
+Note: batch_size=2 is required. ClassificationModule.on_before_batch_transfer
+uses squeeze() on labels; with batch_size=1 this collapses [B] to 0-dim,
+causing CrossEntropyLoss to fail with "batch_size (1) vs (0)".
+"""
+from asparagus.modules.data_modules.training import ClsRegDataModule
+from asparagus.modules.lightning_modules import ClassificationModule
+from asparagus.modules.networks.unet import unet_clsreg_tiny
+from asparagus.pipeline.auto_configuration.checkpoint import load_checkpoint_state_dict
+
+
+def test_test_cls_inference(cls_probe_files, tmp_path, make_trainer):
+    """ClassificationModule runs test-time inference from a saved checkpoint."""
+    ckpt_path = tmp_path / "cls_checkpoint.ckpt"
+
+    # --- Stage 1: train and save a checkpoint ---
+    train_model = unet_clsreg_tiny(input_channels=1, output_channels=2, dimensions="3D")
+    train_module = ClassificationModule(
+        model=train_model,
+        learning_rate=1e-3,
+        warmup_epochs=0,
+        test_output_path=str(tmp_path / "train_preds.json"),
+    )
+    train_dm = ClsRegDataModule(
+        batch_size=2,
+        num_workers=2,  # val_dataloader uses num_workers//2; needs >=2
+        train_split=cls_probe_files["train"],
+        val_split=cls_probe_files["val"],
+        use_random_datasampler=False,
+    )
+    train_trainer = make_trainer()
+    train_trainer.fit(train_module, datamodule=train_dm)
+    train_trainer.save_checkpoint(str(ckpt_path))
+
+    # --- Stage 2: load weights and run inference (mirrors test_cls.py logic) ---
+    weights = load_checkpoint_state_dict(str(ckpt_path))
+    infer_model = unet_clsreg_tiny(input_channels=1, output_channels=2, dimensions="3D")
+    infer_module = ClassificationModule(
+        model=infer_model,
+        weights=weights,
+        test_output_path=str(tmp_path / "test_preds.json"),
+    )
+    test_dm = ClsRegDataModule(
+        batch_size=2,
+        num_workers=2,  # val_dataloader uses num_workers//2; needs >=2
+        train_split=None,
+        val_split=None,
+        test_samples=cls_probe_files["test"],
+        use_random_datasampler=False,
+    )
+    make_trainer(limit_test_batches=2).test(infer_module, datamodule=test_dm)

tests/test_train_reg.py

@@ -0,0 +1,33 @@
+"""Integration test for pipeline/run/train_reg.py components.
+
+Uses RegressionModule + ClsRegDataModule + unet_clsreg_tiny on synthetic (image, label) data.
+Runs both trainer.fit() and trainer.test() mirroring the full pipeline.
+"""
+from asparagus.modules.data_modules.training import ClsRegDataModule
+from asparagus.modules.lightning_modules import RegressionModule
+from asparagus.modules.networks.unet import unet_clsreg_tiny
+
+
+def test_train_reg_fit_and_test(reg_files, tmp_path, make_trainer):
+    """RegressionModule fits then runs inference on synthetic data."""
+    model = unet_clsreg_tiny(input_channels=1, output_channels=1, dimensions="3D")
+
+    data_module = ClsRegDataModule(
+        batch_size=1,
+        num_workers=2,  # val_dataloader uses num_workers//2; needs >=2
+        train_split=reg_files["train"],
+        val_split=reg_files["val"],
+        test_samples=reg_files["test"],
+        use_random_datasampler=False,
+    )
+
+    module = RegressionModule(
+        model=model,
+        learning_rate=1e-3,
+        warmup_epochs=0,
+        test_output_path=str(tmp_path / "preds.json"),
+    )
+
+    trainer = make_trainer(limit_test_batches=1)
+    trainer.fit(module, datamodule=data_module)
+    trainer.test(module, datamodule=data_module)