update test_plot.py

luochang212 · May 24, 2024 · 0a2506a · 0a2506a
1 parent 893379f
commit 0a2506a
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/tests/test_plot.py b/tests/test_plot.py
@@ -1,13 +1,17 @@
+from unittest import mock
+
 from flameai.plot import confusion_matrix, roc_curve
 
 
 def test_roc_curve() -> None:
     y_true = [0, 1, 1, 0, 1, 1, 0, 0, 1, 1]
     y_score = [0.1, 0.4, 0.35, 0.8, 0.15, 0.35, 0.2, 0.7, 0.05, 0.9]
-    roc_curve(y_true, y_score)
+    with mock.patch('matplotlib.pyplot.show'):
+        roc_curve(y_true, y_score)
 
 
 def test_confusion_matrix() -> None:
     y_true = [0, 1, 1, 0, 1, 1, 0, 0, 0, 1]
     y_label = [1, 1, 1, 0, 1, 0, 0, 0, 1, 1]
-    confusion_matrix(y_true, y_label)
+    with mock.patch('matplotlib.pyplot.show'):
+        confusion_matrix(y_true, y_label)