Fix bugs + add tests

sdv-dev · Feb 6, 2023 · cb52e4d · cb52e4d
1 parent 875c39f
commit cb52e4d
Show file tree

Hide file tree

Showing 4 changed files with 53 additions and 4 deletions.
diff --git a/sdv/single_table/copulagan.py b/sdv/single_table/copulagan.py
@@ -198,8 +198,8 @@ def _fit(self, processed_data):
         unseen_columns = self._numerical_distributions.keys() - set(processed_data.columns)
         for column in unseen_columns:
             LOGGER.info(
-                f"Requested distribution {self.numerical_distributions['column']} "
-                f'cannot be applied to column {column} because it no longer '
+                f"Requested distribution '{self.numerical_distributions[column]}' "
+                f"cannot be applied to column '{column}' because it no longer "
                 'exists after preprocessing.'
             )
 

diff --git a/sdv/single_table/copulas.py b/sdv/single_table/copulas.py
@@ -127,8 +127,8 @@ def _fit(self, processed_data):
         unseen_columns = self._numerical_distributions.keys() - set(processed_data.columns)
         for column in unseen_columns:
             LOGGER.info(
-                f"Requested distribution {self.numerical_distributions['column']} "
-                f'cannot be applied to column {column} because it no longer '
+                f"Requested distribution '{self.numerical_distributions[column]}' "
+                f"cannot be applied to column '{column}' because it no longer "
                 'exists after preprocessing.'
             )
 

diff --git a/tests/unit/single_table/test_copulagan.py b/tests/unit/single_table/test_copulagan.py
@@ -224,6 +224,31 @@ def test__create_gaussian_normalizer_config(self, mock_rdt):
         assert config == expected_config
         assert mock_rdt.transformers.GaussianNormalizer.call_args_list == expected_calls
 
+    @patch('sdv.single_table.copulagan.LOGGER')
+    @patch('sdv.single_table.copulagan.CTGANSynthesizer._fit')
+    @patch('sdv.single_table.copulagan.rdt')
+    def test__fit_logging(self, mock_rdt, mock_ctgansynthesizer__fit, mock_logger):
+        """Test a message is logged.
+
+        A message should be logged if the columns passed in ``numerical_distributions``
+        were renamed/dropped during preprocessing.
+        """
+        # Setup
+        metadata = SingleTableMetadata()
+        metadata.add_column('col', sdtype='numerical')
+        numerical_distributions = {'col': 'gamma'}
+        instance = CopulaGANSynthesizer(metadata, numerical_distributions=numerical_distributions)
+        processed_data = pd.DataFrame()
+
+        # Run
+        instance._fit(processed_data)
+
+        # Assert
+        mock_logger.info.assert_called_once_with(
+            "Requested distribution 'gamma' cannot be applied to column 'col' "
+            'because it no longer exists after preprocessing.'
+        )
+
     @patch('sdv.single_table.copulagan.CTGANSynthesizer._fit')
     @patch('sdv.single_table.copulagan.rdt')
     def test__fit(self, mock_rdt, mock_ctgansynthesizer__fit):

diff --git a/tests/unit/single_table/test_copulas.py b/tests/unit/single_table/test_copulas.py
@@ -120,6 +120,30 @@ def test_get_parameters(self):
             'default_distribution': 'beta'
         }
 
+    @patch('sdv.single_table.copulas.LOGGER')
+    def test__fit_logging(self, mock_logger):
+        """Test a message is logged.
+
+        A message should be logged if the columns passed in ``numerical_distributions``
+        were renamed/dropped during preprocessing.
+        """
+        # Setup
+        metadata = SingleTableMetadata()
+        metadata.add_column('col', sdtype='numerical')
+        numerical_distributions = {'col': 'gamma'}
+        instance = GaussianCopulaSynthesizer(
+            metadata, numerical_distributions=numerical_distributions)
+        processed_data = pd.DataFrame({'updated_col': [1, 2, 3]})
+
+        # Run
+        instance._fit(processed_data)
+
+        # Assert
+        mock_logger.info.assert_called_once_with(
+            "Requested distribution 'gamma' cannot be applied to column 'col' "
+            'because it no longer exists after preprocessing.'
+        )
+
     @patch('sdv.single_table.copulas.warnings')
     @patch('sdv.single_table.copulas.multivariate')
     def test__fit(self, mock_multivariate, mock_warnings):