diff --git a/sdmetrics/reports/multi_table/_properties/base.py b/sdmetrics/reports/multi_table/_properties/base.py index a6cd0a72..36fb6e2a 100644 --- a/sdmetrics/reports/multi_table/_properties/base.py +++ b/sdmetrics/reports/multi_table/_properties/base.py @@ -56,9 +56,13 @@ def _get_num_iterations(self, metadata): def _extract_tuple(data, relation): parent_data = data[relation['parent_table_name']] child_data = data[relation['child_table_name']] + primary_key = relation['parent_primary_key'] + foreign_key = relation['child_foreign_key'] + primary_key = [primary_key] if not isinstance(primary_key, list) else primary_key + foreign_key = [foreign_key] if not isinstance(foreign_key, list) else foreign_key return ( - parent_data[relation['parent_primary_key']], - child_data[relation['child_foreign_key']], + parent_data[primary_key], + child_data[foreign_key], ) def _compute_average(self): diff --git a/tests/unit/reports/multi_table/_properties/test_base.py b/tests/unit/reports/multi_table/_properties/test_base.py index 27dbb5c8..df9b54bb 100644 --- a/tests/unit/reports/multi_table/_properties/test_base.py +++ b/tests/unit/reports/multi_table/_properties/test_base.py @@ -120,7 +120,8 @@ def test__extract_tuple(self): real_columns = base_property._extract_tuple(real_data, relation) # Assert - assert real_columns == (real_data['users']['user_id'], real_data['sessions']['user_id']) + pd.testing.assert_frame_equal(real_columns[0], real_data['users'][['user_id']]) + pd.testing.assert_frame_equal(real_columns[1], real_data['sessions'][['user_id']]) def test__generate_details_property(self): """Test the ``_generate_details`` method.""" diff --git a/tests/unit/reports/multi_table/_properties/test_relationship_validity.py b/tests/unit/reports/multi_table/_properties/test_relationship_validity.py index 815d2c8e..8edd448c 100644 --- a/tests/unit/reports/multi_table/_properties/test_relationship_validity.py +++ b/tests/unit/reports/multi_table/_properties/test_relationship_validity.py @@ -92,7 +92,31 @@ def test__extract_tuple(self, real_data_fixture): real_columns = relationship_validity._extract_tuple(real_data, relation) # Assert - assert real_columns == (real_data['users']['user_id'], real_data['sessions']['user_id']) + pd.testing.assert_frame_equal(real_columns[0], real_data['users'][['user_id']]) + pd.testing.assert_frame_equal(real_columns[1], real_data['sessions'][['user_id']]) + + def test__generate_details_referential_integrity_score( + self, real_data_fixture, synthetic_data_fixture, metadata_fixture + ): + """Test that ``ReferentialIntegrity`` returns a score instead of erroring. + + ``_extract_tuple`` used to pass the key columns as ``pd.Series``, which made + ``ReferentialIntegrity`` fail with ``AttributeError: 'Series' object has no + attribute 'columns'`` and record a ``nan`` score. + """ + # Setup + relationship_validity = RelationshipValidity() + + # Run + relationship_validity._generate_details( + real_data_fixture, synthetic_data_fixture, metadata_fixture + ) + + # Assert + details = relationship_validity.details + ri_row = details[details['Metric'] == 'ReferentialIntegrity'].iloc[0] + assert ri_row['Error'] is None + assert not np.isnan(ri_row['Score']) def test__get_num_iteration(self): """Test the ``_get_num_iterations`` method."""