--- a
+++ b/test/unit_tests/test_util.py
@@ -0,0 +1,369 @@
+# Authors: Hubert Banville <hubert.jbanville@gmail.com>
+#          Bruno Aristimunha <b.aristimunha@gmail.com>
+# License: BSD-3
+
+import os
+from unittest import mock
+
+import mne
+import numpy as np
+import h5py
+import pytest
+import torch
+import tempfile
+
+from sklearn.utils import check_random_state
+from numpy.testing import assert_array_equal, assert_allclose
+
+
+from braindecode.util import _cov_and_var_to_corr, _cov_to_corr, \
+    corr, create_mne_dummy_raw, \
+    read_all_file_names, set_random_seeds, th_to_np, cov, np_to_th, \
+    get_balanced_batches
+
+
+def test_create_mne_dummy_raw(tmp_path):
+    n_channels, n_times, sfreq = 2, 10000, 100
+    raw, fnames = create_mne_dummy_raw(
+        n_channels, n_times, sfreq, savedir=tmp_path, save_format=["fif", "hdf5"]
+    )
+
+    assert isinstance(raw, mne.io.RawArray)
+    assert len(raw.ch_names) == n_channels
+    assert raw.n_times == n_times
+    assert raw.info["sfreq"] == sfreq
+    assert isinstance(fnames, dict)
+    assert os.path.isfile(fnames["fif"])
+    assert os.path.isfile(fnames["hdf5"])
+
+    _ = mne.io.read_raw_fif(fnames["fif"], preload=False, verbose=None)
+    with h5py.File(fnames["hdf5"], "r") as hf:
+        _ = np.array(hf["fake_raw"])
+
+
+def test_set_random_seeds_raise_value_error():
+    with pytest.raises(
+        ValueError, match="cudnn_benchmark expected to be bool or None, got 'abc'"
+    ):
+        set_random_seeds(100, True, "abc")
+
+
+def test_set_random_seeds_warning():
+    torch.backends.cudnn.benchmark = True
+    with pytest.warns(
+        UserWarning,
+        match="torch.backends.cudnn.benchmark was set to True which may results in "
+        "lack of reproducibility. In some cases to ensure reproducibility you "
+        "may need to set torch.backends.cudnn.benchmark to False.",
+    ):
+        set_random_seeds(100, True)
+
+
+def test_set_random_seeds_with_valid_cudnn_benchmark():
+    with mock.patch("torch.backends.cudnn") as mock_cudnn:
+        # Test with cudnn_benchmark = True
+        set_random_seeds(42, cuda=True, cudnn_benchmark=True)
+        assert mock_cudnn.benchmark is True
+
+        # Test with cudnn_benchmark = False
+        set_random_seeds(42, cuda=True,
+                         cudnn_benchmark=False)
+        assert mock_cudnn.benchmark is False
+
+
+def test_set_random_seeds_with_invalid_cudnn_benchmark():
+    with pytest.raises(ValueError):
+        set_random_seeds(42, cuda=True,
+                         cudnn_benchmark='invalid_type')
+
+
+def test_th_to_np_data_preservation():
+    # Test with different data types
+    for dtype in [torch.float32, torch.int32]:
+        tensor = torch.tensor([1, 2, 3], dtype=dtype)
+        np_array = th_to_np(tensor)
+        assert np_array.dtype == tensor.numpy().dtype
+        # Corrected attribute access
+        assert_array_equal(np_array, tensor.numpy())
+
+
+def test_th_to_np_on_cpu():
+    # Create a tensor on CPU
+    cpu_tensor = torch.tensor([1, 2, 3], dtype=torch.float32)
+    np_array = th_to_np(cpu_tensor)
+    # Check the type and data of the numpy array
+    assert isinstance(np_array, np.ndarray)
+    assert np_array.dtype == cpu_tensor.numpy().dtype
+    # Correct way to check dtype
+    assert_array_equal(np_array, cpu_tensor.numpy())
+
+
+def test_cov_basic():
+    # Create two simple identical arrays
+    a = np.array([[1, 2, 3], [4, 5, 6]])
+    b = np.array([[1, 2, 3], [4, 5, 6]])
+    expected_cov = np.array([[1, 1], [1, 1]])  # Calculated expected covariance
+    computed_cov = cov(a, b)
+    assert_allclose(computed_cov, expected_cov, rtol=1e-5)
+
+
+def test_cov_dimension_mismatch():
+    # Arrays with mismatched sample size should raise an error
+    a = np.array([[1, 2], [3, 4]])
+    b = np.array([[1, 2, 3]])
+    with pytest.raises(ValueError):
+        cov(a, b)
+
+
+def test_np_to_th_basic_conversion():
+    # Convert a simple list to tensor
+    data = [1, 2, 3]
+    tensor = np_to_th(data)
+    assert torch.is_tensor(tensor)
+    assert_array_equal(tensor.numpy(), np.array(data))
+
+
+def test_np_to_th_dtype_conversion():
+    # Convert and specify dtype
+    data = [1.0, 2.0, 3.0]
+    tensor = np_to_th(data, dtype=np.float32)
+    assert tensor.dtype == torch.float32
+
+
+def test_np_to_th_requires_grad():
+    # Check requires_grad attribute
+    data = np.array([1.0, 2.0, 3.0])
+    tensor = np_to_th(data, requires_grad=True)
+    assert tensor.requires_grad is True
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(),
+                    reason="Requires CUDA support")
+def test_np_to_th_pin_memory():
+    # Create a numpy array
+    data = np.array([1, 2, 3])
+
+    # Convert the numpy array to a tensor with pin_memory=True
+    tensor = np_to_th(data, pin_memory=True)
+    # Check if the tensor is pinned in memory
+    assert tensor.is_pinned() is True
+    # Convert the numpy array to a tensor with pin_memory=False
+    tensor = np_to_th(data, pin_memory=False)
+    # Check if the tensor is not pinned in memory
+    assert tensor.is_pinned() is False
+
+def test_np_to_th_requires_grad_unsupported_dtype():
+    # Attempt to set requires_grad on an unsupported dtype (integers)
+    data = np.array([1, 2, 3])
+    with pytest.raises(RuntimeError,
+                       match="Only Tensors of floating point and complex dtype can require gradients"):
+        np_to_th(data, requires_grad=True)
+
+
+def test_np_to_th_tensor_options():
+    # Additional tensor options like device
+    data = [1, 2, 3]
+    device = 'cuda' if torch.cuda.is_available() else 'cpu'
+    tensor = np_to_th(data, device=device)
+    assert tensor.device.type == device
+
+
+def test_np_to_th_single_number_conversion():
+    # Single number conversion to tensor
+    data = 42
+    tensor = np_to_th(data)
+    assert torch.is_tensor(tensor)
+    assert tensor.item() == 42  # Check if the value is correct
+
+
+def test_cov_and_var_to_corr_zero_variance():
+    # Scenario with zero variance, expecting result to handle divide by zero
+    this_cov = np.array([[1, 0],
+                         [0, 1]])
+    var_a = np.array([0, 1])  # Zero variance in the first variable
+    var_b = np.array([1, 0])  # Zero variance in the second variable
+    calculated_corr = _cov_and_var_to_corr(this_cov, var_a, var_b)
+
+    # Expected correlation matrix
+    expected_corr = np.array([[np.inf, np.nan],
+                               [0, np.inf]])
+    assert_array_equal(calculated_corr, expected_corr)
+
+
+def test_cov_and_var_to_corr_single_element():
+    # Testing with single-element arrays
+    this_cov = np.array([[1]])
+    var_a = np.array([1])
+    var_b = np.array([1])
+    expected_corr = np.array([[1]])
+    calculated_corr = _cov_and_var_to_corr(this_cov, var_a, var_b)
+    assert_array_equal(calculated_corr, expected_corr)
+
+
+
+def test_cov_to_corr_unbiased():
+    # Create datasets a and b with known covariance and variance characteristics
+    a = np.array([[1, 2, 3, 4],
+                  [2, 3, 4, 5]])
+    b = np.array([[1, 3, 5, 7],
+                  [5, 6, 7, 8]])
+    # Covariance between the features of a and b
+    # Calculating covariance manually for known values
+    demeaned_a = a - np.mean(a, axis=1, keepdims=True)
+    demeaned_b = b - np.mean(b, axis=1, keepdims=True)
+    this_cov = np.dot(demeaned_a, demeaned_b.T) / (b.shape[1] - 1)
+
+    # Compute expected correlation using standard formulas for correlation
+    var_a = np.var(a, axis=1, ddof=1)
+    var_b = np.var(b, axis=1, ddof=1)
+    expected_divisor = np.outer(np.sqrt(var_a), np.sqrt(var_b))
+    expected_corr = this_cov / expected_divisor
+
+    # Compute correlation using the function
+    calculated_corr = _cov_to_corr(this_cov, a, b)
+
+    # Assert that the calculated correlation matches the expected correlation
+    assert_allclose(calculated_corr, expected_corr, rtol=1e-5)
+
+
+def test_balanced_batches_basic():
+    n_trials = 100
+    seed = 42
+    rng = check_random_state(seed)
+    n_batches = 10
+    batches = get_balanced_batches(n_trials, rng, shuffle=False,
+                                   n_batches=n_batches)
+
+    # Check correct number of batches
+    assert len(batches) == n_batches
+
+    # Check balanced batch sizes
+    all_batch_sizes = [len(batch) for batch in batches]
+    max_size = max(all_batch_sizes)
+    min_size = min(all_batch_sizes)
+    assert max_size - min_size <= 1
+
+    # Check if all indices are unique and accounted for
+    all_indices = np.concatenate(batches)
+    assert np.array_equal(np.sort(all_indices), np.arange(n_trials))
+
+
+def test_balanced_batches_with_batch_size():
+    n_trials = 105
+    seed = 42
+    rng = check_random_state(seed)
+    batch_size = 20
+    batches = get_balanced_batches(n_trials, rng, shuffle=False,
+                                   batch_size=batch_size)
+
+    # Check the modified batch size condition
+    expected_n_batches = int(np.round(n_trials / float(batch_size)))
+    assert len(batches) == expected_n_batches
+
+    # Checking the total number of indices
+    all_indices = np.concatenate(batches)
+    assert len(all_indices) == n_trials
+
+
+def test_balanced_batches_shuffle():
+    n_trials = 50
+    seed = 42
+    rng = check_random_state(seed)
+    batches_no_shuffle = get_balanced_batches(n_trials, rng, shuffle=False,
+                                              batch_size=10)
+    rng = check_random_state(seed)
+    batches_with_shuffle = get_balanced_batches(n_trials, rng, shuffle=True,
+                                                batch_size=10)
+
+    # Check that shuffling changes the order of indices
+    assert not np.array_equal(np.concatenate(batches_no_shuffle),
+                              np.concatenate(batches_with_shuffle))
+
+
+def test_corr_correlation_computation():
+    # Create two 2D arrays with known correlation
+    a = np.array([[1, 2, 3], [4, 5, 6]])
+    b = np.array([[1, 2, 3], [4, 5, 6]])
+
+    # Call the corr function
+    corr_result = corr(a, b)
+
+    # Compute the known correlation
+    known_corr = np.array([np.corrcoef(a[i], b[i]) for i in range(a.shape[0])])
+
+    # Extract the correlation computation from the corr function
+    computed_corr = _cov_to_corr(cov(a, b), a, b)
+
+    # Assert that the computed correlation matches the known correlation
+    assert np.allclose(computed_corr, known_corr)
+    assert np.allclose(corr_result, computed_corr)
+
+
+def test_get_balanced_batches_zero_batches():
+    # Create a scenario where n_batches is 0
+    n_trials = 10
+    rng = check_random_state(0)
+    shuffle = False
+    n_batches = 0
+    batch_size = None
+
+    # Call the get_balanced_batches function
+    batches = get_balanced_batches(n_trials, rng, shuffle, n_batches, batch_size)
+
+    # Check if the function returns a single batch with all trials
+    assert len(batches) == 1
+    assert len(batches[0]) == n_trials
+
+
+def test_get_balanced_batches_i_batch_less_than_n_batches_with_extra_trial():
+    # Create a scenario where i_batch < n_batches_with_extra_trial
+    n_trials = 10
+    rng = check_random_state(0)
+    shuffle = False
+    n_batches = 6
+    batch_size = None
+
+    # Call the get_balanced_batches function
+    batches = get_balanced_batches(n_trials, rng, shuffle, n_batches, batch_size)
+
+    # Check if the first batch has one more trial than the last batch
+    assert len(batches[0]) > len(batches[-1])
+
+
+def test_read_all_file_names():
+    # Create a temporary directory
+    with tempfile.TemporaryDirectory() as tmpdir:
+        temp_files = []
+        try:
+            # Create some temporary files with .txt extension
+            for i in range(5):
+                temp_file = os.path.join(tmpdir, f'temp{i}.txt')
+                with open(temp_file, 'w') as f:
+                    f.write('This is a temporary file.')
+                temp_files.append(temp_file)
+
+            # Call the read_all_file_names function
+            file_paths = read_all_file_names(tmpdir, '.txt')
+
+            # Check if the function found all the temporary files
+            assert len(file_paths) == 5
+
+            # Check if the paths returned by the function are correct
+            for i in range(5):
+                assert os.path.join(tmpdir, f'temp{i}.txt') in file_paths
+        finally:
+            # Delete the temporary files
+            for temp_file in temp_files:
+                os.remove(temp_file)
+
+
+def test_read_all_file_names_error():
+    with pytest.raises(AssertionError):
+        # Call the read_all_file_names function with a non-existent directory
+        read_all_file_names('non_existent_dir', '.txt')
+
+
+def test_read_all_files_not_extension():
+    with pytest.raises(AssertionError):
+        # Call the read_all_file_names function with a non-existent directory
+        read_all_file_names('non_existent_dir', 'txt')