_validate for multi table report

sdv-dev · Nov 20, 2023 · 3bb2392 · 3bb2392
1 parent c716f20
commit 3bb2392
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 5 deletions.
diff --git a/sdmetrics/reports/base_report.py b/sdmetrics/reports/base_report.py
@@ -81,10 +81,6 @@ def _validate(self, real_data, synthetic_data, metadata):
         """
         self._validate_data_format(real_data, synthetic_data)
         if self.__class__.__name__ == 'DiagnosticReport':
-            table_name = list(metadata.get('tables', {}).keys())
-            if table_name:
-                self.table_names = table_name
-
             return
 
         self._validate_metadata_matches_data(real_data, synthetic_data, metadata)

diff --git a/sdmetrics/reports/multi_table/base_multi_table_report.py b/sdmetrics/reports/multi_table/base_multi_table_report.py
@@ -60,14 +60,27 @@ def _validate_relationships(self, real_data, synthetic_data, metadata):
 
     def _validate_metadata_matches_data(self, real_data, synthetic_data, metadata):
         """Validate that the metadata matches the data."""
-        self.table_names = list(metadata['tables'].keys())
         for table in self.table_names:
             super()._validate_metadata_matches_data(
                 real_data[table], synthetic_data[table], metadata['tables'][table]
             )
 
         self._validate_relationships(real_data, synthetic_data, metadata)
 
+    def _validate(self, real_data, synthetic_data, metadata):
+        """Validate the inputs.
+
+        Args:
+            real_data (pandas.DataFrame):
+                The real data.
+            synthetic_data (pandas.DataFrame):
+                The synthetic data.
+            metadata (dict):
+                The metadata of the table.
+        """
+        self.table_names = list(metadata['tables'].keys())
+        super()._validate(real_data, synthetic_data, metadata)
+
     def _check_table_names(self, table_name):
         if table_name not in self.table_names:
             raise ValueError(f"Unknown table ('{table_name}'). Must be one of {self.table_names}.")

diff --git a/tests/unit/reports/multi_table/test_base_multi_table_report.py b/tests/unit/reports/multi_table/test_base_multi_table_report.py
@@ -131,6 +131,7 @@ def test__validate_metadata_matches_data(self, mock__validate_metadata_matches_d
 
         report = BaseMultiTableReport()
         report._validate_relationships = mock__validate_relationships
+        report.table_names = ['Table_1', 'Table_2']
 
         # Run
         report._validate_metadata_matches_data(real_data, synthetic_data, metadata)