def test_columns_multiple_columns(self):
"""Test column names mapped correctly for multiple columns."""
dummy_data = pd.DataFrame({"col1": [1, 2, 3], "col2": [2, 3, 4]})
test_masker = Masker()
dummy_data = test_masker.fit_transform(dummy_data)
assert all(dummy_data.columns == ["column_0", "column_1"])
def test_columns_single_column(self):
"""Test column names mapped correctly for a single column."""
dummy_data = pd.DataFrame({"col1": [1, 2, 3]})
test_masker = Masker()
dummy_data = test_masker.fit_transform(dummy_data)
assert dummy_data.columns[0] == "column_0"
def test_values_single_column(self):
"""Test column values mapped correctly for a single column."""
dummy_data = pd.DataFrame({"col1": ["A", "B", "A"]})
test_masker = Masker()
dummy_data = test_masker.fit_transform(dummy_data)
assert all(dummy_data.iloc[:, 0] == ["level_0", "level_1", "level_0"])
def test_values_multiple_columns(self):
"""Test column values maped correctly for multple columns."""
dummy_data = pd.DataFrame({
"col1": ["A", "B", "C"],
"col2": ["1", "1", "1"]
})
test_masker = Masker()
dummy_data = test_masker.fit_transform(dummy_data)
assert all(dummy_data.iloc[:, 0] == ["level_0", "level_1", "level_2"])
assert all(dummy_data.iloc[:, 1] == ["level_0", "level_0", "level_0"])
def test_dataframe_matches_with_scaling(self):
"""Full test that the dataframe is as expected, with numeric scaling."""
dummy_data = pd.DataFrame({
"col1": ["1", "2", "1"],
"col2": [1, 2, 3],
"col3": ["red", "red", "blue"]
})
expected_data = pd.DataFrame({
"column_0": ["level_0", "level_1", "level_0"],
"column_1": [0.0, 0.5, 1.0],
"column_2": ["level_0", "level_0", "level_1"],
})
test_masker = Masker(numerical_scaling=True)
dummy_data = test_masker.fit_transform(dummy_data)
assert dummy_data.equals(expected_data)
