Update test_plot_confusion_matrix.py

avigupta2612 · web-flow · commit 0c296abf5228 · 2020-06-26T16:46:27.000+05:30
diff --git a/sklearn/metrics/_plot/tests/test_plot_confusion_matrix.py b/sklearn/metrics/_plot/tests/test_plot_confusion_matrix.py
@@ -249,6 +249,17 @@ def test_confusion_matrix_pipeline(pyplot, clf, data, n_classes):
     assert disp.text_.shape == (n_classes, n_classes)
 
 
+@pytest.mark.parametrize("colorbar", [True, False])
+def test_plot_confusion_matrix_colorbar(pyplot, data, fitted_clf, colorbar):
+    X, y = data
+    disp = plot_confusion_matrix(fitted_clf, X, y, colorbar=True)
+    assert disp.im_.colorbar is not None
+
+    # attempt a plot with the opposit effect
+    disp.plot(colorbar=False)
+    assert disp.im_.colorbar is None
+
+
 @pytest.mark.parametrize("values_format", ['e', 'n'])
 def test_confusion_matrix_text_format(pyplot, data, y_pred, n_classes,
                                       fitted_clf, values_format):
