tests

Author: R-Palazzo

tests/integration/synthesizer/test_ctgan.py

@@ -236,7 +236,7 @@ def test_fixed_random_seed():
     })
     discrete_columns = ['discrete']
 
-    ctgan = CTGAN(epochs=1, cuda=False)
+    ctgan = CTGAN(epochs=1, enable_gpu=False)
 
     # Run
     ctgan.fit(data, discrete_columns)

tests/unit/synthesizer/test__utils.py

@@ -0,0 +1,78 @@
+import re
+import sys
+from unittest.mock import patch
+
+import pytest
+import torch
+
+from ctgan.synthesizers._utils import _set_device, _validate_gpu_parameters, validate_and_set_device
+
+
+def test__validate_gpu_parameter():
+    """Test the ``_validate_gpu_parameters`` method."""
+    # Setup
+    expected_error = re.escape(
+        'Cannot resolve the provided values of `cuda` and `enable_gpu` parameters. '
+        'Please use only `enable_gpu`.'
+    )
+    expected_warning = re.escape(
+        '`cuda` parameter is deprecated and will be removed in a future release. '
+        'Please use `enable_gpu` instead.'
+    )
+
+    # Run
+    enable_gpu_1 = _validate_gpu_parameters(False, None)
+    enable_gpu_2 = _validate_gpu_parameters(True, None)
+    with pytest.warns(FutureWarning, match=expected_warning):
+        enable_gpu_3 = _validate_gpu_parameters(True, False)
+
+    with pytest.raises(ValueError, match=expected_error):
+        _validate_gpu_parameters(False, True)
+
+    # Assert
+    assert enable_gpu_1 is False
+    assert enable_gpu_2 is True
+    assert enable_gpu_3 is False
+
+
+def test__set_device():
+    """Test the ``_set_device`` method."""
+    # Run
+    device_1 = _set_device(False)
+    device_2 = _set_device(True)
+    device_3 = _set_device(True, 'cpu')
+    device_4 = _set_device(enable_gpu=False, device='cpu')
+
+    # Assert
+    if (
+        sys.platform == 'darwin'
+        and getattr(torch.backends, 'mps', None)
+        and torch.backends.mps.is_available()
+    ):
+        expected_device_2 = torch.device('mps')
+    elif torch.cuda.is_available():
+        expected_device_2 = torch.device('cuda')
+    else:
+        expected_device_2 = torch.device('cpu')
+
+    assert device_1 == torch.device('cpu')
+    assert device_2 == expected_device_2
+    assert device_3 == torch.device('cpu')
+    assert device_4 == torch.device('cpu')
+
+
+@patch('ctgan.synthesizers._utils._set_device')
+@patch('ctgan.synthesizers._utils._validate_gpu_parameters')
+def test_validate_and_set_device(mock_validate, mock_set_device):
+    """Test the ``validate_and_set_device`` method."""
+    # Setup
+    mock_validate.return_value = True
+    mock_set_device.return_value = torch.device('cuda')
+
+    # Run
+    device = validate_and_set_device(enable_gpu=True, cuda=None)
+
+    # Assert
+    mock_validate.assert_called_once_with(True, None)
+    mock_set_device.assert_called_once_with(True)
+    assert device == torch.device('cuda')

tests/unit/synthesizer/test_ctgan.py

@@ -1,7 +1,7 @@
 """CTGAN unit testing module."""
 
 from unittest import TestCase
-from unittest.mock import Mock
+from unittest.mock import Mock, patch
 
 import numpy as np
 import pandas as pd
@@ -175,6 +175,30 @@ def _assert_is_between(data, lower, upper):
 
 
 class TestCTGAN(TestCase):
+    @patch('ctgan.synthesizers.ctgan.validate_and_set_device')
+    def test___init__(self, mock_validate_and_set_device):
+        """Test the `__init__` method."""
+        # Setup
+        mock_validate_and_set_device.return_value = 'cpu'
+
+        # Run
+        synth = CTGAN()
+
+        # Assert
+        assert synth._embedding_dim == 128
+        assert synth._generator_dim == (256, 256)
+        assert synth._discriminator_dim == (256, 256)
+        assert synth._batch_size == 500
+        assert synth._epochs == 300
+        assert synth.pac == 10
+        assert synth.loss_values is None
+        assert synth._generator is None
+        assert synth._data_sampler is None
+        assert synth._verbose is False
+        assert synth._enable_gpu is True
+        assert synth._device == 'cpu'
+        mock_validate_and_set_device.assert_called_once_with(True, None)
+
     def test__apply_activate_(self):
         """Test `_apply_activate` for tables with both continuous and categoricals.
 

tests/unit/synthesizer/test_tvae.py

@@ -8,6 +8,27 @@
 
 
 class TestTVAE:
+    @patch('ctgan.synthesizers.tvae.validate_and_set_device')
+    def test___init__(self, mock_validate_and_set_device):
+        """Test the `__init__` method."""
+        # Setup
+        mock_validate_and_set_device.return_value = 'cpu'
+
+        # Run
+        synth = TVAE()
+
+        # Assert
+        assert synth.embedding_dim == 128
+        assert synth.compress_dims == (128, 128)
+        assert synth.decompress_dims == (128, 128)
+        assert synth.batch_size == 500
+        assert synth.epochs == 300
+        assert synth.loss_values.equals(pd.DataFrame(columns=['Epoch', 'Batch', 'Loss']))
+        assert synth.verbose is False
+        assert synth._enable_gpu is True
+        assert synth._device == 'cpu'
+        mock_validate_and_set_device.assert_called_once_with(True, None)
+
     @patch('ctgan.synthesizers.tvae._loss_function')
     @patch('ctgan.synthesizers.tvae.tqdm')
     def test_fit_verbose(self, tqdm_mock, loss_func_mock):