TST Add unit tests for _BinaryClassifierCurveDisplayMixin (scikit-learn#31193)

Author: lucyleeow

sklearn/utils/tests/test_plotting.py

@@ -1,12 +1,120 @@
 import numpy as np
 import pytest
 
+from sklearn.linear_model import LogisticRegression
 from sklearn.utils._plotting import (
+    _BinaryClassifierCurveDisplayMixin,
     _despine,
     _interval_max_min_ratio,
     _validate_score_name,
     _validate_style_kwargs,
 )
+from sklearn.utils._response import _get_response_values_binary
+from sklearn.utils._testing import assert_allclose
+
+
+@pytest.mark.parametrize("ax", [None, "Ax"])
+@pytest.mark.parametrize(
+    "name, expected_name_out", [(None, "TestEstimator"), ("CustomName", "CustomName")]
+)
+def test_validate_plot_params(pyplot, ax, name, expected_name_out):
+    """Check `_validate_plot_params` returns the correct values."""
+    display = _BinaryClassifierCurveDisplayMixin()
+    display.estimator_name = "TestEstimator"
+    if ax:
+        _, ax = pyplot.subplots()
+    ax_out, _, name_out = display._validate_plot_params(ax=ax, name=name)
+
+    assert name_out == expected_name_out
+
+    if ax:
+        assert ax == ax_out
+
+
+@pytest.mark.parametrize("pos_label", [None, 0])
+@pytest.mark.parametrize("name", [None, "CustomName"])
+@pytest.mark.parametrize(
+    "response_method", ["auto", "predict_proba", "decision_function"]
+)
+def test_validate_and_get_response_values(pyplot, pos_label, name, response_method):
+    """Check `_validate_and_get_response_values` returns the correct values."""
+    X = np.array([[0, 0], [1, 1], [2, 2], [3, 3]])
+    y = np.array([0, 0, 2, 2])
+    estimator = LogisticRegression().fit(X, y)
+
+    y_pred, pos_label, name_out = (
+        _BinaryClassifierCurveDisplayMixin._validate_and_get_response_values(
+            estimator,
+            X,
+            y,
+            response_method=response_method,
+            pos_label=pos_label,
+            name=name,
+        )
+    )
+
+    expected_y_pred, expected_pos_label = _get_response_values_binary(
+        estimator, X, response_method=response_method, pos_label=pos_label
+    )
+
+    assert_allclose(y_pred, expected_y_pred)
+    assert pos_label == expected_pos_label
+
+    # Check name is handled correctly
+    expected_name = name if name is not None else "LogisticRegression"
+    assert name_out == expected_name
+
+
+@pytest.mark.parametrize(
+    "y_true, error_message",
+    [
+        (np.array([0, 1, 2]), "The target y is not binary."),
+        (np.array([0, 1]), "Found input variables with inconsistent"),
+        (np.array([0, 2, 0, 2]), r"y_true takes value in \{0, 2\} and pos_label"),
+    ],
+)
+def test_validate_from_predictions_params_errors(pyplot, y_true, error_message):
+    """Check `_validate_from_predictions_params` raises the correct errors."""
+    y_pred = np.array([0.1, 0.2, 0.3, 0.4])
+    sample_weight = np.ones(4)
+
+    with pytest.raises(ValueError, match=error_message):
+        _BinaryClassifierCurveDisplayMixin._validate_from_predictions_params(
+            y_true=y_true,
+            y_pred=y_pred,
+            sample_weight=sample_weight,
+            pos_label=None,
+        )
+
+
+@pytest.mark.parametrize("name", [None, "CustomName"])
+@pytest.mark.parametrize(
+    "pos_label, y_true",
+    [
+        (None, np.array([0, 1, 0, 1])),
+        (2, np.array([0, 2, 0, 2])),
+    ],
+)
+def test_validate_from_predictions_params_returns(pyplot, name, pos_label, y_true):
+    """Check `_validate_from_predictions_params` returns the correct values."""
+    y_pred = np.array([0.1, 0.2, 0.3, 0.4])
+    pos_label_out, name_out = (
+        _BinaryClassifierCurveDisplayMixin._validate_from_predictions_params(
+            y_true=y_true,
+            y_pred=y_pred,
+            sample_weight=None,
+            pos_label=pos_label,
+            name=name,
+        )
+    )
+
+    # Check name is handled correctly
+    expected_name = name if name is not None else "Classifier"
+    assert name_out == expected_name
+
+    # Check pos_label is handled correctly
+    expected_pos_label = pos_label if pos_label is not None else 1
+    assert pos_label_out == expected_pos_label
 
 
 def metric():