fix test

Author: R-Palazzo

rdt/transformers/utils.py

@@ -177,10 +177,11 @@ def strings_from_regex(regex, max_repeat=16):
 
 def _fill_nan_with_none_series(data):
     sentinel = object()
-    if isinstance(data.dtype, pd.CategoricalDtype):
+    dtype = data.dtype
+    if isinstance(dtype, pd.CategoricalDtype):
         data = data.cat.add_categories([sentinel])
 
-    return data.fillna(sentinel).replace({sentinel: None})
+    return data.fillna(sentinel).replace({sentinel: None}).astype(dtype)
 
 
 def fill_nan_with_none(data):

tests/unit/transformers/test_utils.py

@@ -181,10 +181,10 @@ def test__fill_nan_with_none_series():
     # Assert
     expected_result = pd.Series([1.0, 2.0, 3.0, None], dtype='object')
     pd.testing.assert_series_equal(result, expected_result)
-    expected_result_categorical = pd.Series(['a', 'b', 'c', 'd', None], dtype='object')
-    pd.testing.assert_series_equal(
-        result_categorical, expected_result_categorical, check_dtype=False
+    expected_result_categorical = pd.Series(
+        pd.Categorical(['a', 'b', 'c', 'd', None], categories=['a', 'b', 'c', 'd'])
     )
+    pd.testing.assert_series_equal(result_categorical, expected_result_categorical)
 
 
 def test_fill_nan_with_none_series():