From 672203d64f2e1f4135533d83aa480a80409ebd2d Mon Sep 17 00:00:00 2001
From: Falko Habel <falko.habel@gmx.de>
Date: Thu, 27 Mar 2025 17:16:17 +0100
Subject: [PATCH] added test files

---
 tests/finetune/test_trainer.py    | 75 +++++++++++++++++++++++
 tests/inference/test_inference.py | 99 +++++++++++++++++++++++++++++++
 tests/upsampler/test_aiunn.py     | 48 +++++++++++++++
 tests/upsampler/test_config.py    | 45 ++++++++++++++
 4 files changed, 267 insertions(+)
 create mode 100644 tests/finetune/test_trainer.py
 create mode 100644 tests/inference/test_inference.py
 create mode 100644 tests/upsampler/test_aiunn.py
 create mode 100644 tests/upsampler/test_config.py

diff --git a/tests/finetune/test_trainer.py b/tests/finetune/test_trainer.py
new file mode 100644
index 0000000..773f973
--- /dev/null
+++ b/tests/finetune/test_trainer.py
@@ -0,0 +1,75 @@
+import pytest
+import torch
+import torch.nn as nn
+from aiunn import aiuNNTrainer
+
+# Simple mock dataset
+class MockDataset(torch.utils.data.Dataset):
+    def __init__(self, num_samples=10):
+        self.num_samples = num_samples
+    
+    def __len__(self):
+        return self.num_samples
+    
+    def __getitem__(self, idx):
+        return torch.randn(3, 64, 64), torch.randn(3, 128, 128)
+
+# Simple mock model
+class MockModel(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.conv = nn.Conv2d(3, 3, 3, padding=1)
+    
+    def forward(self, x):
+        return self.conv(x)
+    
+    def save(self, path):
+        pass  # Mock save method
+
+@pytest.fixture
+def trainer():
+    model = MockModel()
+    return aiuNNTrainer(model, dataset_class=MockDataset)
+
+def test_trainer_initialization(trainer):
+    """Test basic trainer initialization"""
+    assert trainer.model is not None
+    assert isinstance(trainer.criterion, nn.MSELoss)
+    assert trainer.optimizer is None
+    assert trainer.device in [torch.device('cuda'), torch.device('cpu')]
+
+def test_load_data_basic(trainer):
+    """Test basic data loading"""
+    train_loader, val_loader = trainer.load_data(
+        dataset_params={'num_samples': 10},
+        batch_size=2,
+        validation_split=0.2
+    )
+    
+    assert train_loader is not None
+    assert val_loader is not None
+    assert len(train_loader) > 0
+    assert len(val_loader) > 0
+
+def test_load_custom_datasets(trainer):
+    """Test loading custom datasets"""
+    train_dataset = MockDataset(num_samples=10)
+    val_dataset = MockDataset(num_samples=5)
+    
+    train_loader, val_loader = trainer.load_data(
+        custom_train_dataset=train_dataset,
+        custom_val_dataset=val_dataset,
+        batch_size=2
+    )
+    
+    assert train_loader is not None
+    assert val_loader is not None
+    assert len(train_loader) == 5  # 10 samples with batch size 2
+    assert len(val_loader) == 3    # 5 samples with batch size 2 (rounded up)
+
+def test_error_no_dataset():
+    """Test error when no dataset is provided"""
+    trainer = aiuNNTrainer(MockModel(), dataset_class=None)
+    
+    with pytest.raises(ValueError):
+        trainer.load_data(dataset_params={})
\ No newline at end of file
diff --git a/tests/inference/test_inference.py b/tests/inference/test_inference.py
new file mode 100644
index 0000000..39dea9a
--- /dev/null
+++ b/tests/inference/test_inference.py
@@ -0,0 +1,99 @@
+# tests/inference/test_inference.py
+import pytest
+import numpy as np
+import torch
+from PIL import Image
+from aiunn import aiuNNInference
+from aiunn.upsampler.aiunn import aiuNN, aiuNNConfig
+from aiia import AIIABase, AIIAConfig
+import os
+import json
+from unittest.mock import patch, MagicMock, mock_open
+
+@pytest.fixture
+def real_model(tmp_path):
+    # Create temporary directory for model
+    model_dir = tmp_path / "model"
+    model_dir.mkdir()
+    
+    config = AIIAConfig()
+    ai_config = aiuNNConfig()
+    base_model = AIIABase(config)
+    
+    # Make sure aiuNN is properly configured with all required attributes
+    upsampler = aiuNN(base_model, config=ai_config)
+    # Ensure the upsample attribute is properly set if needed
+    # upsampler.upsample = ... # Add any necessary initialization
+    
+    # Save the model and config to temporary directory
+    save_path = str(model_dir / "save")
+    os.makedirs(save_path, exist_ok=True)
+    
+    # Save config file
+    config_data = {
+        "model_type": "test_model",
+        "scale": 4,
+        "in_channels": 3,
+        "out_channels": 3
+    }
+    with open(os.path.join(save_path, "config.json"), "w") as f:
+        json.dump(config_data, f)
+    
+    # Save model
+    upsampler.save(save_path)
+    
+    # Load model in inference mode
+    inference_model = aiuNNInference(model_path=save_path, precision='fp16', device='cpu')
+    return inference_model
+
+
+@pytest.fixture
+def inference(real_model):
+    return real_model
+
+def test_preprocess_image(inference):
+    # Create a small test image
+    test_array = np.zeros((100, 100, 3), dtype=np.uint8)
+    test_image = Image.fromarray(test_array)
+    
+    # Test with PIL Image
+    result = inference.preprocess_image(test_image)
+    assert isinstance(result, torch.Tensor)
+    assert result.shape[0] == 1  # batch dimension
+    assert result.shape[1] == 3  # channels
+    assert result.shape[2:] == (100, 100)  # height, width
+
+def test_postprocess_tensor(inference):
+    # Create a test tensor
+    test_tensor = torch.zeros(1, 3, 100, 100)
+    
+    result = inference.postprocess_tensor(test_tensor)
+    assert isinstance(result, Image.Image)
+    assert result.size == (100, 100)
+    assert result.mode == 'RGB'
+
+def test_save(inference):
+    # Create a test image
+    test_image = Image.fromarray(np.zeros((100, 100, 3), dtype=np.uint8))
+    output_path = "test_output.png"
+    
+    with patch('os.makedirs') as mock_makedirs:
+        inference.save(test_image, output_path)
+        mock_makedirs.assert_called_with(os.path.dirname(os.path.abspath(output_path)), exist_ok=True)
+
+def test_convert_to_binary(inference):
+    # Create a test image
+    test_image = Image.fromarray(np.zeros((100, 100, 3), dtype=np.uint8))
+    
+    result = inference.convert_to_binary(test_image)
+    assert isinstance(result, bytes)
+    assert len(result) > 0
+
+def test_process_batch(inference):
+    # Create test images
+    test_array = np.zeros((100, 100, 3), dtype=np.uint8)
+    test_images = [Image.fromarray(test_array) for _ in range(2)]
+    
+    results = inference.process_batch(test_images)
+    assert len(results) == 2
+    assert all(isinstance(img, Image.Image) for img in results)
\ No newline at end of file
diff --git a/tests/upsampler/test_aiunn.py b/tests/upsampler/test_aiunn.py
new file mode 100644
index 0000000..aae0813
--- /dev/null
+++ b/tests/upsampler/test_aiunn.py
@@ -0,0 +1,48 @@
+import os
+import tempfile
+from aiia import AIIABase, AIIAConfig
+from aiunn import aiuNN, aiuNNConfig
+
+def test_save_and_load_model():
+    # Create a temporary directory to save the model
+    with tempfile.TemporaryDirectory() as tmpdirname:
+        # Create configurations and build a base model
+        config = AIIAConfig()
+        ai_config = aiuNNConfig()
+        base_model = AIIABase(config)
+        upsampler = aiuNN(base_model, config=ai_config)
+
+        # Save the model
+        save_path = os.path.join(tmpdirname, "model")
+        upsampler.save(save_path)
+
+        # Load the model
+        loaded_upsampler = aiuNN.load(save_path)
+
+        # Verify that the loaded model is the same as the original model
+        assert isinstance(loaded_upsampler, aiuNN)
+        assert loaded_upsampler.config.__dict__ == upsampler.config.__dict__
+
+def test_save_and_load_model_with_precision():
+    # Create a temporary directory to save the model
+    with tempfile.TemporaryDirectory() as tmpdirname:
+        # Create configurations and build a base model
+        config = AIIAConfig()
+        ai_config = aiuNNConfig()
+        base_model = AIIABase(config)
+        upsampler = aiuNN(base_model, config=ai_config)
+
+        # Save the model
+        save_path = os.path.join(tmpdirname, "model")
+        upsampler.save(save_path)
+
+        # Load the model with precision 'bf16'
+        loaded_upsampler = aiuNN.load(save_path, precision="bf16")
+
+        # Verify that the loaded model is the same as the original model
+        assert isinstance(loaded_upsampler, aiuNN)
+        assert loaded_upsampler.config.__dict__ == upsampler.config.__dict__
+
+if __name__ == "__main__":
+    test_save_and_load_model()
+    test_save_and_load_model_with_precision()
\ No newline at end of file
diff --git a/tests/upsampler/test_config.py b/tests/upsampler/test_config.py
new file mode 100644
index 0000000..3eb2779
--- /dev/null
+++ b/tests/upsampler/test_config.py
@@ -0,0 +1,45 @@
+import pytest
+from aiunn import aiuNNConfig
+
+def test_default_initialization():
+    config = aiuNNConfig()
+    assert config.upsample_scale == 2
+    assert config.upsample_mode == 'bilinear'
+    assert not config.upsample_align_corners
+    assert len(config.layers) == 1
+    assert config.layers[0]['name'] == 'Upsample'
+
+def test_custom_initialization():
+    custom_config = {
+        'some_key': 'some_value',
+    }
+    config = aiuNNConfig(base_config=custom_config, upsample_scale=3, upsample_mode='nearest', upsample_align_corners=True)
+    assert config.upsample_scale == 3
+    assert config.upsample_mode == 'nearest'
+    assert config.upsample_align_corners
+    assert len(config.layers) == 1
+    assert config.layers[0]['name'] == 'Upsample'
+
+def test_add_upsample_layer():
+    config = aiuNNConfig()
+    config.add_upsample_layer()
+    assert len(config.layers) == 1
+
+def test_upsample_layer_not_duplicated():
+    config = aiuNNConfig()
+    initial_length = len(config.layers)
+    # Remove all existing 'Upsample' layers.
+    for layer in list(config.layers):
+        if layer.get('name') == 'Upsample':
+            config.layers.remove(layer)
+    config.add_upsample_layer()
+    assert len(config.layers) == 1
+
+def test_base_config_with_to_dict():
+    class MockBaseConfig:
+        def to_dict(self):
+            return {'base_key': 'base_value'}
+
+    base_config = MockBaseConfig()
+    config = aiuNNConfig(base_config=base_config)
+    assert config.to_dict()['base_key'] == 'base_value'
\ No newline at end of file