Repository: scikit-learn/scikit-learn
File: sklearn/inspection/_plot/tests/test_boundary_decision_display.py (Python)

Code before:
def test_input_validation_errors(pyplot, kwargs, error_msg, fitted_clf):
    """Check input validation from_estimator."""
    with pytest.raises(ValueError, match=error_msg):
        DecisionBoundaryDisplay.from_estimator(fitted_clf, X, **kwargs)


@pytest.mark.parametrize(
    "kwargs, error_msg",
    [
        (
            {"multiclass_colors": {"dict": "not_list"}},
            "'multiclass_colors' must be a list or a str.",
        ),
        ({"multiclass_colors": "not_cmap"}, "it must be a valid Matplotlib colormap"),
        ({"multiclass_colors": ["red", "green"]}, "it must be of the same length"),
        (
            {"multiclass_colors": ["red", "green", "blue", "yellow"]},
            "it must be of the same length",
        ),
        (
            {"multiclass_colors": ["red", "green", "not color"]},
            "it can only contain valid Matplotlib color names",
        ),
    ],
)
def test_input_validation_errors_multiclass_colors(pyplot, kwargs, error_msg):
    """Check input validation for `multiclass_colors` in `from_estimator`."""
    X, y = load_iris_2d_scaled()
    clf = LogisticRegression().fit(X, y)
    with pytest.raises(ValueError, match=error_msg):
        DecisionBoundaryDisplay.from_estimator(clf, X, **kwargs)


def test_display_plot_input_error(pyplot, fitted_clf):
    """Check input validation for `plot`."""
    disp = DecisionBoundaryDisplay.from_estimator(fitted_clf, X, grid_resolution=5)

    with pytest.raises(ValueError, match="plot_method must be 'contourf'"):
        disp.plot(plot_method="hello_world")


@pytest.mark.parametrize(
    "response_method", ["auto", "predict", "predict_proba", "decision_function"]
)
@pytest.mark.parametrize("plot_method", ["contourf", "contour"])
def test_decision_boundary_display_classifier(
    pyplot, fitted_clf, response_method, plot_method
):
    """Check that decision boundary is correct."""
    fig, ax = pyplot.subplots()
    eps = 2.0

Review comment (bug):
Added edge case for too many values (could also be accepted and select the first n elements, but currently we check for exact length)

Code after:
def test_input_validation_errors(pyplot, kwargs, error_msg, fitted_clf):
    """Check input validation from_estimator."""
    with pytest.raises(ValueError, match=error_msg):
        DecisionBoundaryDisplay.from_estimator(fitted_clf, X, **kwargs)


@pytest.mark.parametrize(
    "kwargs, error_type, error_msg",
    [
        (
            {"multiclass_colors": {"dict": "not_list"}},
            TypeError,
            "'multiclass_colors' must be a list or a str.",
        ),
        (
            {"multiclass_colors": "not_cmap"},
            ValueError,
            "it must be a valid Matplotlib colormap",
        ),
        (
            {"multiclass_colors": ["red", "green"]},
            ValueError,
            "it must be of the same length",
        ),
        (
            {"multiclass_colors": ["red", "green", "blue", "yellow"]},
            ValueError,
            "it must be of the same length",
        ),
        (
            {"multiclass_colors": ["red", "green", "not color"]},
            ValueError,
            "it can only contain valid Matplotlib color names",
        ),
    ],
)
def test_input_validation_errors_multiclass_colors(
    pyplot, kwargs, error_type, error_msg
):
    """Check input validation for `multiclass_colors` in `from_estimator`."""
    X, y = load_iris_2d_scaled()
    clf = LogisticRegression().fit(X, y)
    with pytest.raises(error_type, match=error_msg):
        DecisionBoundaryDisplay.from_estimator(clf, X, **kwargs)


def test_display_plot_input_error(pyplot, fitted_clf):
    """Check input validation for `plot`."""
    disp = DecisionBoundaryDisplay.from_estimator(fitted_clf, X, grid_resolution=5)

    with pytest.raises(ValueError, match="plot_method must be 'contourf'"):