From da309d003d1afcdd66d4968d1713dfb6ec340248 Mon Sep 17 00:00:00 2001 From: R-Palazzo Date: Fri, 23 Aug 2024 15:53:59 +0200 Subject: [PATCH] integration test --- tests/integration/constraints/test_tabular.py | 37 +++++++++++++++++++ 1 file changed, 37 insertions(+) diff --git a/tests/integration/constraints/test_tabular.py b/tests/integration/constraints/test_tabular.py index 23fa701f0..32bb6e7d2 100644 --- a/tests/integration/constraints/test_tabular.py +++ b/tests/integration/constraints/test_tabular.py @@ -70,3 +70,40 @@ def test_fixed_combinations_with_nans(): data.drop_duplicates(ignore_index=True), check_like=True, ) + + +def test_fixedincrements_with_nullable_pandas_dtypes(): + """Test that FixedIncrements constraint works with nullable pandas dtypes.""" + # Setup + data = pd.DataFrame({ + 'UInt8': pd.Series([1, pd.NA, 3], dtype='UInt8') * 10, + 'UInt16': pd.Series([1, pd.NA, 4], dtype='UInt16') * 10, + 'UInt32': pd.Series([1, pd.NA, 5], dtype='UInt32') * 10, + 'UInt64': pd.Series([1, pd.NA, 6], dtype='UInt64') * 10, + }) + metadata = SingleTableMetadata().load_from_dict({ + 'columns': { + 'UInt8': {'sdtype': 'numerical', 'computer_representation': 'UInt8'}, + 'UInt16': {'sdtype': 'numerical', 'computer_representation': 'UInt16'}, + 'UInt32': {'sdtype': 'numerical', 'computer_representation': 'UInt32'}, + 'UInt64': {'sdtype': 'numerical', 'computer_representation': 'UInt64'}, + } + }) + gcs = GaussianCopulaSynthesizer(metadata) + my_constraints = [ + { + 'constraint_class': 'FixedIncrements', + 'constraint_parameters': {'column_name': column, 'increment_value': 10}, + } + for column in data.columns + ] + gcs.add_constraints(my_constraints) + + # Run + gcs.fit(data) + synthetic_data = gcs.sample(10) + + # Assert + synthetic_data.dtypes.to_dict() == data.dtypes.to_dict() + for column in data.columns: + assert np.all(synthetic_data[column] % 10 == 0)