Skip to content

Commit

Permalink
add integration test integer columns only
Browse files Browse the repository at this point in the history
  • Loading branch information
R-Palazzo committed Aug 16, 2024
1 parent ba22a99 commit 617d5c4
Showing 1 changed file with 33 additions and 0 deletions.
33 changes: 33 additions & 0 deletions tests/integration/constraints/test_tabular.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,39 @@
from sdv.single_table import GaussianCopulaSynthesizer


def test_fixed_combinations_integers():
"""Test that FixedCombinations constraint works with integer columns."""
data = pd.DataFrame({
'A': [1, 2, 3, 1, 2, 1],
'B': [10, 20, 30, 10, 20, 10],
})
metadata = SingleTableMetadata().load_from_dict({
'columns': {
'A': {'sdtype': 'categorical'},
'B': {'sdtype': 'categorical'},
}
})

synthesizer = GaussianCopulaSynthesizer(metadata)
my_constraint = {
'constraint_class': 'FixedCombinations',
'constraint_parameters': {'column_names': ['A', 'B']},
}
synthesizer.add_constraints(constraints=[my_constraint])

# Run
synthesizer.fit(data)
synthetic_data = synthesizer.sample(1000)

# Assert
assert len(synthetic_data) == 1000
pd.testing.assert_frame_equal(
synthetic_data.drop_duplicates(ignore_index=True),
data.drop_duplicates(ignore_index=True),
check_like=True,
)


def test_fixed_combinations_with_nans():
"""Test that FixedCombinations constraint works with NaNs."""
# Setup
Expand Down

0 comments on commit 617d5c4

Please sign in to comment.