add tests

sdv-dev · Oct 24, 2024 · b2e5087 · b2e5087
1 parent ea2e92b
commit b2e5087
Showing 1 changed file with 26 additions and 10 deletions.
diff --git a/tests/unit/data_processing/test_data_processor.py b/tests/unit/data_processing/test_data_processor.py
@@ -1137,7 +1137,9 @@ def test__create_config(self):
             'first_name': ['John', 'Doe', 'Johanna'],
             'id': ['ID_001', 'ID_002', 'ID_003'],
             'id_no_regex': ['ID_001', 'ID_002', 'ID_003'],
-            'id_numeric': [0, 1, 2],
+            'id_numeric_int8': pd.Series([1, 2, 3], dtype='Int8'),
+            'id_numeric_int16': pd.Series([1, 2, 3], dtype='Int16'),
+            'id_numeric_int32': pd.Series([1, 2, 3], dtype='Int32'),
             'id_column': ['ID_999', 'ID_999', 'ID_007'],
             'date': ['2021-02-01', '2022-03-05', '2023-01-31'],
             'unknown': ['a', 'b', 'c'],
@@ -1151,9 +1153,9 @@ def test__create_config(self):
         dp.create_anonymized_transformer.return_value = 'AnonymizedFaker'
         dp.create_regex_generator.return_value = 'RegexGenerator'
         dp.metadata.primary_key = 'id'
-        dp.metadata.alternate_keys = ['id_no_regex', 'id_numeric']
+        dp.metadata.alternate_keys = ['id_no_regex', 'id_numeric_int8']
         dp._primary_key = 'id'
-        dp._keys = ['id', 'id_no_regex', 'id_numeric']
+        dp._keys = ['id', 'id_no_regex', 'id_numeric_int8']
         dp.metadata.columns = {
             'int': {'sdtype': 'numerical'},
             'float': {'sdtype': 'numerical'},
@@ -1163,7 +1165,9 @@ def test__create_config(self):
             'first_name': {'sdtype': 'first_name'},
             'id': {'sdtype': 'id', 'regex_format': 'ID_\\d{3}[0-9]'},
             'id_no_regex': {'sdtype': 'id'},
-            'id_numeric': {'sdtype': 'id'},
+            'id_numeric_int8': {'sdtype': 'id'},
+            'id_numeric_int16': {'sdtype': 'id'},
+            'id_numeric_int32': {'sdtype': 'id'},
             'id_column': {'sdtype': 'id'},
             'date': {'sdtype': 'datetime', 'datetime_format': '%Y-%m-%d'},
             'unknown': {'sdtype': 'unknown'},
@@ -1188,7 +1192,9 @@ def test__create_config(self):
             'first_name': 'pii',
             'id': 'text',
             'id_no_regex': 'text',
-            'id_numeric': 'text',
+            'id_numeric_int8': 'text',
+            'id_numeric_int16': 'text',
+            'id_numeric_int32': 'text',
             'id_column': 'text',
             'date': 'datetime',
             'unknown': 'pii',
@@ -1236,11 +1242,21 @@ def test__create_config(self):
         assert id_no_regex_transformer.function_kwargs == {'text': 'sdv-id-??????'}
         assert id_no_regex_transformer.cardinality_rule == 'unique'
 
-        id_numeric_transformer = config['transformers']['id_numeric']
-        assert isinstance(id_numeric_transformer, AnonymizedFaker)
-        assert id_numeric_transformer.function_name == 'bothify'
-        assert id_numeric_transformer.function_kwargs == {'text': '#########'}
-        assert id_numeric_transformer.cardinality_rule == 'unique'
+        id_numeric_int_8_transformer = config['transformers']['id_numeric_int8']
+        assert isinstance(id_numeric_int_8_transformer, AnonymizedFaker)
+        assert id_numeric_int_8_transformer.function_name == 'bothify'
+        assert id_numeric_int_8_transformer.function_kwargs == {'text': '##'}
+        assert id_numeric_int_8_transformer.cardinality_rule == 'unique'
+
+        id_numeric_int_16_transformer = config['transformers']['id_numeric_int16']
+        assert isinstance(id_numeric_int_16_transformer, AnonymizedFaker)
+        assert id_numeric_int_16_transformer.function_name == 'bothify'
+        assert id_numeric_int_16_transformer.function_kwargs == {'text': '####'}
+
+        id_numeric_int_32_transformer = config['transformers']['id_numeric_int32']
+        assert isinstance(id_numeric_int_32_transformer, AnonymizedFaker)
+        assert id_numeric_int_32_transformer.function_name == 'bothify'
+        assert id_numeric_int_32_transformer.function_kwargs == {'text': '#########'}
 
         id_column_transformer = config['transformers']['id_column']
         assert isinstance(id_column_transformer, AnonymizedFaker)