def test_categorical_map_attribute_created(self, data):
"""Test that the categorical_map attribute is created for the Masker object."""
test_masker = Masker()
test_masker.get_categorical_map(X=data)
assert hasattr(test_masker, "categorical_map")
def test_X_type(self):
"""Test error is thrown if X is not the correct type."""
test_masker = Masker()
with pytest.raises(TypeError):
test_masker.get_categorical_map(X=123)
def test_categorical_map_value(self, data):
"""Test that the categorical_map attribute takes the correct value."""
test_masker = Masker()
test_masker.get_categorical_map(X=data)
expected_map = {
"col1": {
"one": "level_0",
"two": "level_1",
"three": "level_2"
},
"col3": {
"1": "level_0",
"2": "level_1"
},
"col4": {
np.NaN: "level_0",
"blue": "level_1"
},
}
assert test_masker.categorical_map == expected_map