add int tests with multiple FKs

Author: gsheni

sdv/multi_table/hma.py

@@ -350,7 +350,8 @@ def _get_extension(self, child_name, child_table, foreign_key, progress_bar_desc
         extension_rows = []
         primary_key = self.metadata.tables[child_name].primary_key
         foreign_key_columns = self.metadata._get_all_foreign_keys(child_name)
-        if primary_key and primary_key in foreign_key_columns:
+        primary_key_is_a_foreign_key = primary_key and primary_key in foreign_key_columns
+        if primary_key_is_a_foreign_key and foreign_key == primary_key:
             # data processor will set index of each table to the PK for table
             foreign_key_values = child_table.index.unique()
         else:

tests/integration/multi_table/conftest.py

@@ -1,3 +1,4 @@
+import random
 from copy import deepcopy
 
 import pandas as pd
@@ -173,3 +174,91 @@ def data_metadata_1_to_1_to_1_subset_to_subset(data_metadata_1_to_1_subset_diamo
     ]
     metadata = Metadata.load_from_dict(metadata_dict)
     return data, metadata
+
+
+@pytest.fixture
+def data_metadata_multiple_foreign_keys():
+    parent_1_ids = range(0, 10)
+    parent_2_ids = range(10, 20)
+    parent = pd.DataFrame({
+        'parent_id': parent_1_ids,
+        'col_categorical': random.choices(['A', 'B', 'C', 'D', 'E'], k=10),
+    })
+    child = pd.DataFrame({
+        'parent_1_id': parent_1_ids,
+        'parent_2_id': parent_2_ids,
+        'col_numerical': [10.2, 20.3] * 5,
+    })
+    second_parent = pd.DataFrame({'parent_id': parent_2_ids, 'col_boolean': [True, False] * 5})
+    data = {
+        'parent': parent,
+        'child': child,
+        'second_parent': second_parent,
+    }
+    metadata = Metadata.load_from_dict({
+        'tables': {
+            'parent': {
+                'columns': {
+                    'parent_id': {'sdtype': 'id'},
+                    'col_categorical': {'sdtype': 'categorical'},
+                },
+                'primary_key': 'parent_id',
+            },
+            'child': {
+                'columns': {
+                    'parent_1_id': {'sdtype': 'id'},
+                    'parent_2_id': {'sdtype': 'id'},
+                    'col_numerical': {'sdtype': 'numerical'},
+                },
+                'primary_key': 'parent_1_id',
+            },
+            'second_parent': {
+                'columns': {'parent_id': {'sdtype': 'id'}, 'col_boolean': {'sdtype': 'boolean'}},
+                'primary_key': 'parent_id',
+            },
+        },
+        'relationships': [
+            {
+                'parent_table_name': 'parent',
+                'child_table_name': 'child',
+                'parent_primary_key': 'parent_id',
+                'child_foreign_key': 'parent_1_id',
+            },
+            {
+                'parent_table_name': 'second_parent',
+                'child_table_name': 'child',
+                'parent_primary_key': 'parent_id',
+                'child_foreign_key': 'parent_2_id',
+            },
+        ],
+    })
+    assert data['child']['parent_1_id'].equals(data['parent']['parent_id'])
+    assert data['child']['parent_2_id'].equals(data['second_parent']['parent_id'])
+    metadata.validate()
+    metadata.validate_data(data)
+    return data, metadata
+
+
+@pytest.fixture
+def data_metadata_multiple_foreign_keys_subset(data_metadata_multiple_foreign_keys):
+    _, metadata = data_metadata_multiple_foreign_keys
+    parent = pd.DataFrame({'parent_id': [1, 2, 3], 'col_categorical': ['A', 'B', 'C']})
+    child = pd.DataFrame({
+        'parent_1_id': [1, 2],
+        'parent_2_id': [1, 1],
+        'col_numerical': [100.5, 101.6],
+    })
+    second_parent = pd.DataFrame({
+        'parent_id': [1, 2],
+        'col_boolean': [True, False],
+    })
+    data = {
+        'parent': parent,
+        'child': child,
+        'second_parent': second_parent,
+    }
+    assert set(data['child']['parent_1_id']).issubset(set(data['parent']['parent_id']))
+    assert set(data['child']['parent_2_id']).issubset(set(data['second_parent']['parent_id']))
+    metadata.validate()
+    metadata.validate_data(data)
+    return data, metadata

tests/integration/multi_table/test_hma.py

@@ -2858,95 +2858,128 @@ def test_datetime_warning_doesnt_repeat():
     assert len(matching_warnings) == 1
 
 
-def test_hma_1_to_1(data_metadata_1_to_1):
-    """Test HMA handles PK to PK relationship (1 to 1) and synthetic data matching cardinality."""
-    # Setup
-    data, metadata = data_metadata_1_to_1
+class TestPrimaryKeyToPrimaryKey:
+    def test_1_to_1(self, data_metadata_1_to_1):
+        """Test HMA handles PK to PK relationship (1 to 1) and synthetic data matching cardinality."""
+        # Setup
+        data, metadata = data_metadata_1_to_1
 
-    # Run
-    synthesizer = HMASynthesizer(metadata=metadata, verbose=False)
-    with warnings.catch_warnings(record=True) as caught_warnings:
+        # Run
+        synthesizer = HMASynthesizer(metadata=metadata, verbose=False)
+        with warnings.catch_warnings(record=True) as caught_warnings:
+            synthesizer.fit(data)
+            synthetic_data = synthesizer.sample(scale=1)
+
+        # Assert
+        assert synthetic_data['guests']['guest_email'].equals(
+            synthetic_data['rooms']['guest_email']
+        )
+        synthesizer.validate(synthetic_data)
+        for msg in caught_warnings:
+            assert 'ChainedAssignmentError' not in str(msg.message)
+
+    def test_1_to_1_or_0(self, data_metadata_1_to_1_or_0):
+        """Test HMA handles PK to PK relationship (1 to 1/0) and synthetic data matching cardinality."""
+        # Setup
+        data, metadata = data_metadata_1_to_1_or_0
+
+        # Run
+        synthesizer = HMASynthesizer(metadata=metadata, verbose=False)
         synthesizer.fit(data)
         synthetic_data = synthesizer.sample(scale=1)
 
-    # Assert
-    assert synthetic_data['guests']['guest_email'].equals(synthetic_data['rooms']['guest_email'])
-    synthesizer.validate(synthetic_data)
-    for msg in caught_warnings:
-        assert 'ChainedAssignmentError' not in str(msg.message)
+        # Assert
+        assert set(synthetic_data['users']['user_id']).issuperset(
+            set(synthetic_data['survey_response']['user_id'])
+        )
+        synthesizer.validate(synthetic_data)
 
+    def test_1_to_1_or_0_not_superset(self, data_metadata_1_to_1_or_0):
+        """Test error is raised if PK to PK relationship but parent is not a superset."""
+        # Setup
+        data, metadata = data_metadata_1_to_1_or_0
+        metadata.remove_relationship(parent_table_name='users', child_table_name='survey_response')
+        metadata.add_relationship(
+            parent_table_name='survey_response',
+            parent_primary_key='user_id',
+            child_table_name='users',
+            child_foreign_key='user_id',
+        )
+        synthesizer = HMASynthesizer(metadata=metadata, verbose=False)
+        match_ = re.escape("Error: foreign key column 'user_id' contains unknown references: (9).")
 
-def test_hma_1_to_1_or_0(data_metadata_1_to_1_or_0):
-    """Test HMA handles PK to PK relationship (1 to 1/0) and synthetic data matching cardinality."""
-    # Setup
-    data, metadata = data_metadata_1_to_1_or_0
+        # Run and Assert
+        with pytest.raises(InvalidDataError, match=match_):
+            synthesizer.fit(data)
 
-    # Run
-    synthesizer = HMASynthesizer(metadata=metadata, verbose=False)
-    synthesizer.fit(data)
-    synthetic_data = synthesizer.sample(scale=1)
+    def test_1_to_1_to_1_subset_to_subset(self, data_metadata_1_to_1_to_1_subset_to_subset):
+        """Test PK to PK to PK, with the 2nd and 3rd table having a subset."""
+        # Setup
+        data, metadata = data_metadata_1_to_1_to_1_subset_to_subset
 
-    # Assert
-    assert set(synthetic_data['users']['user_id']).issuperset(
-        set(synthetic_data['survey_response']['user_id'])
-    )
-    synthesizer.validate(synthetic_data)
+        # Run
+        synthesizer = HMASynthesizer(metadata=metadata, verbose=False)
+        synthesizer.fit(data)
+        synthetic_data = synthesizer.sample(scale=1.0)
 
+        # Assert
+        assert set(synthetic_data['guests']['guest_email']).issuperset(
+            set(synthetic_data['guests_pii']['guest_email'])
+        )
+        assert set(synthetic_data['guests_pii']['guest_email']).issuperset(
+            set(synthetic_data['rooms']['guest_email'])
+        )
+        synthesizer.validate(synthetic_data)
 
-def test_hma_1_to_1_or_0_not_superset(data_metadata_1_to_1_or_0):
-    """Test error is raised if PK to PK relationship but parent is not a superset."""
-    # Setup
-    data, metadata = data_metadata_1_to_1_or_0
-    metadata.remove_relationship(parent_table_name='users', child_table_name='survey_response')
-    metadata.add_relationship(
-        parent_table_name='survey_response',
-        parent_primary_key='user_id',
-        child_table_name='users',
-        child_foreign_key='user_id',
-    )
-    synthesizer = HMASynthesizer(metadata=metadata, verbose=False)
-    match_ = re.escape("Error: foreign key column 'user_id' contains unknown references: (9).")
+    def test_1_to_1_to_1_diamond(self, data_metadata_1_to_1_subset_diamond):
+        """Test PK to PK to PK in a diamond relationship."""
+        # Setup
+        data, metadata = data_metadata_1_to_1_subset_diamond
 
-    # Run and Assert
-    with pytest.raises(InvalidDataError, match=match_):
+        # Run
+        synthesizer = HMASynthesizer(metadata=metadata, verbose=False)
         synthesizer.fit(data)
+        synthetic_data = synthesizer.sample(scale=1.0)
 
+        # Assert
+        assert set(synthetic_data['guests']['guest_email']).issuperset(
+            set(synthetic_data['guests_pii']['guest_email'])
+        )
+        assert set(synthetic_data['guests']['guest_email']).issuperset(
+            set(synthetic_data['rooms']['guest_email'])
+        )
+        synthesizer.validate(synthetic_data)
 
-def test_1_to_1_to_1_subset_to_subset(data_metadata_1_to_1_to_1_subset_to_subset):
-    """Test PK to PK to PK, with the 2nd and 3rd table having a subset."""
-    # Setup
-    data, metadata = data_metadata_1_to_1_to_1_subset_to_subset
-
-    # Run
-    synthesizer = HMASynthesizer(metadata=metadata, verbose=False)
-    synthesizer.fit(data)
-    synthetic_data = synthesizer.sample(scale=1.0)
+    def test_multiple_fks(self, data_metadata_multiple_foreign_keys):
+        """Test support for parent and child with multiple foreign keys."""
+        # Setup
+        data, metadata = data_metadata_multiple_foreign_keys
 
-    # Assert
-    assert set(synthetic_data['guests']['guest_email']).issuperset(
-        set(synthetic_data['guests_pii']['guest_email'])
-    )
-    assert set(synthetic_data['guests_pii']['guest_email']).issuperset(
-        set(synthetic_data['rooms']['guest_email'])
-    )
-    synthesizer.validate(synthetic_data)
+        # Run
+        synthesizer = HMASynthesizer(metadata=metadata, verbose=False)
+        synthesizer.fit(data)
+        synthetic_data = synthesizer.sample(scale=1.0)
 
+        # Assert
+        for each_parent_id in synthetic_data['child']['parent_1_id'].tolist():
+            assert each_parent_id in set(synthetic_data['parent']['parent_id'])
+        for each_parent_id in synthetic_data['child']['parent_2_id'].tolist():
+            assert each_parent_id in set(synthetic_data['second_parent']['parent_id'])
+        synthesizer.validate(synthetic_data)
 
-def test_1_to_1_to_1_diamond(data_metadata_1_to_1_subset_diamond):
-    """Test PK to PK to PK in a diamond relationship."""
-    # Setup
-    data, metadata = data_metadata_1_to_1_subset_diamond
+    def test_multiple_fks_mismatched(self, data_metadata_multiple_foreign_keys_mismatched):
+        """Test support for parent and child with multiple foreign keys."""
+        # Setup
+        data, metadata = data_metadata_multiple_foreign_keys_mismatched
 
-    # Run
-    synthesizer = HMASynthesizer(metadata=metadata, verbose=False)
-    synthesizer.fit(data)
-    synthetic_data = synthesizer.sample(scale=1.0)
+        # Run
+        synthesizer = HMASynthesizer(metadata=metadata, verbose=False)
+        synthesizer.fit(data)
+        synthetic_data = synthesizer.sample(scale=1.0)
 
-    # Assert
-    assert set(synthetic_data['guests']['guest_email']).issuperset(
-        set(synthetic_data['guests_pii']['guest_email'])
-    )
-    assert set(synthetic_data['guests']['guest_email']).issuperset(
-        set(synthetic_data['rooms']['guest_email'])
-    )
-    synthesizer.validate(synthetic_data)
+        # Assert
+        for each_parent_id in synthetic_data['child']['parent_1_id'].tolist():
+            assert each_parent_id in set(synthetic_data['parent']['parent_id'])
+        for each_parent_id in synthetic_data['child']['parent_2_id'].tolist():
+            assert each_parent_id in set(synthetic_data['second_parent']['parent_id'])
+        synthesizer.validate(synthetic_data)