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

Code before:
            7,
            2,
        ),
    ],
)
def test_n_classes_attribute(pyplot, estimator, n_blobs, expected_n_classes):
    """Check that `n_classes` is set correctly.

    Introduced in https://github.com/scikit-learn/scikit-learn/pull/33015"""

    X, y = make_blobs(n_samples=150, centers=n_blobs, n_features=2, random_state=42)
    clf = estimator.fit(X, y)
    disp = DecisionBoundaryDisplay.from_estimator(clf, X, response_method="predict")
    assert disp.n_classes == expected_n_classes

    # Test that setting class_of_interest always converts to a binary problem.
    disp_coi = DecisionBoundaryDisplay.from_estimator(
        clf, X, class_of_interest=y[0], response_method="predict"
    )
    assert disp_coi.n_classes == 2


def test_n_classes_raises_if_not_inferrable(pyplot):
    """Check behaviour if `n_classes` can't be inferred.

    Non-regression test for issue #33194"""

    class CustomUnknownEstimator(BaseEstimator):
        def fit(self, X, y):
            self.fitted_ = True
            return self

        def predict(self, X):
            return np.array(0)

    X, y = load_iris_2d_scaled()
    est = CustomUnknownEstimator().fit(X, y)
    msg = "Number of classes or labels cannot be inferred from CustomUnknownEstimator"
    with pytest.raises(ValueError, match=msg):
        DecisionBoundaryDisplay.from_estimator(est, X, response_method="predict")


def test_cmap_and_colors_logic(pyplot):
    """Check the handling logic for `cmap` and `colors`."""
    X, y = load_iris_2d_scaled()
    clf = LogisticRegression().fit(X, y)

    with pytest.warns(
        UserWarning,
        match="'cmap' is ignored in favor of 'multiclass_colors'",
    ):

Review comment (suggestion):
```suggestion
    Non-regression test for issue #33194.
    """
```
nit

Code after:
            7,
            2,
        ),
    ],
)
def test_n_classes_attribute(pyplot, estimator, n_blobs, expected_n_classes):
    """Check that `n_classes` is set correctly.

    Introduced in https://github.com/scikit-learn/scikit-learn/pull/33015.
    """

    X, y = make_blobs(n_samples=150, centers=n_blobs, n_features=2, random_state=42)
    clf = estimator.fit(X, y)
    disp = DecisionBoundaryDisplay.from_estimator(clf, X, response_method="predict")
    assert disp.n_classes == expected_n_classes

    # Test that setting class_of_interest always converts to a binary problem.
    disp_coi = DecisionBoundaryDisplay.from_estimator(
        clf, X, class_of_interest=y[0], response_method="predict"
    )
    assert disp_coi.n_classes == 2


def test_n_classes_raises_if_not_inferrable(pyplot):
    """Check behaviour if `n_classes` can't be inferred.

    Non-regression test for issue #33194.
    """

    class CustomUnknownEstimator(BaseEstimator):
        def fit(self, X, y):
            self.fitted_ = True
            return self

        def predict(self, X):
            return np.array(0)

    X, y = load_iris_2d_scaled()
    est = CustomUnknownEstimator().fit(X, y)
    msg = "Number of classes or labels cannot be inferred from CustomUnknownEstimator"
    with pytest.raises(ValueError, match=msg):
        DecisionBoundaryDisplay.from_estimator(est, X, response_method="predict")


def test_cmap_and_colors_logic(pyplot):
    """Check the handling logic for `cmap` and `colors`."""
    X, y = load_iris_2d_scaled()
    clf = LogisticRegression().fit(X, y)

    with pytest.warns(
        UserWarning,