Repository: scikit-learn/scikit-learn
File: sklearn/impute/tests/test_impute.py (Python)

Code before:
    with pytest.raises(ValueError, match=re.escape(err_msg)):
        imputer = SimpleImputer(
            missing_values=missing_value, strategy="constant", fill_value=fill_value
        )
        imputer.fit_transform(X)


def test_imputation_constant_integer():
    # Test imputation using the constant strategy on integers
    X = np.array([[-1, 2, 3, -1], [4, -1, 5, -1], [6, 7, -1, -1], [8, 9, 0, -1]])

    X_true = np.array([[0, 2, 3, 0], [4, 0, 5, 0], [6, 7, 0, 0], [8, 9, 0, 0]])

    imputer = SimpleImputer(
        missing_values=-1, strategy="constant", fill_value=0, keep_empty_features=True
    )
    X_trans = imputer.fit_transform(X)

    assert_array_equal(X_trans, X_true)

    imputer = SimpleImputer(
        missing_values=-1, strategy="constant", fill_value=0, keep_empty_features=False
    )
    X_trans = imputer.fit_transform(X)

    assert_array_equal(X_trans, X_true[:, :-1])


@pytest.mark.parametrize("array_constructor", CSR_CONTAINERS + [np.asarray])
def test_imputation_constant_float(array_constructor):
    # Test imputation using the constant strategy on floats
    X = np.array(
        [
            [np.nan, 1.1, 0, np.nan],
            [1.2, np.nan, 1.3, np.nan],
            [0, 0, np.nan, np.nan],
            [1.4, 1.5, 0, np.nan],
        ]
    )

    X_true = np.array(
        [[-1, 1.1, 0, -1], [1.2, -1, 1.3, -1], [0, 0, -1, -1], [1.4, 1.5, 0, -1]]
    )

    X = array_constructor(X)

    X_true = array_constructor(X_true)

    imputer = SimpleImputer(
        strategy="constant", fill_value=-1, keep_empty_features=True
    )

Review comment (suggestion):
Let's parametrize the test instead over `keep_empty_feature` and then you can define `expected_X_trans = X_true if keep_empty_features else X_true[:, :-1]`

Code after:
    with pytest.raises(ValueError, match=re.escape(err_msg)):
        imputer = SimpleImputer(
            missing_values=missing_value, strategy="constant", fill_value=fill_value
        )
        imputer.fit_transform(X)


@pytest.mark.parametrize("keep_empty_features", [True, False])
def test_imputation_constant_integer(keep_empty_features):
    # Test imputation using the constant strategy on integers
    X = np.array([[-1, 2, 3, -1], [4, -1, 5, -1], [6, 7, -1, -1], [8, 9, 0, -1]])

    X_true = np.array([[0, 2, 3, 0], [4, 0, 5, 0], [6, 7, 0, 0], [8, 9, 0, 0]])
    if not keep_empty_features:
        X_true = X_true[:, :-1]

    imputer = SimpleImputer(
        missing_values=-1,
        strategy="constant",
        fill_value=0,
        keep_empty_features=keep_empty_features,
    )
    X_trans = imputer.fit_transform(X)

    assert_array_equal(X_trans, X_true)


@pytest.mark.parametrize("array_constructor", CSR_CONTAINERS + [np.asarray])
@pytest.mark.parametrize("keep_empty_features", [True, False])
def test_imputation_constant_float(array_constructor, keep_empty_features):
    # Test imputation using the constant strategy on floats
    X = np.array(
        [
            [np.nan, 1.1, 0, np.nan],
            [1.2, np.nan, 1.3, np.nan],
            [0, 0, np.nan, np.nan],
            [1.4, 1.5, 0, np.nan],
        ]
    )

    X_true = np.array(
        [[-1, 1.1, 0, -1], [1.2, -1, 1.3, -1], [0, 0, -1, -1], [1.4, 1.5, 0, -1]]
    )
    if not keep_empty_features:
        X_true = X_true[:, :-1]

    X = array_constructor(X)

    X_true = array_constructor(X_true)

    imputer = SimpleImputer(