Use mocker over unittest.mock

tests/contrib/test_viz.py

@@ -1,6 +1,5 @@
 import json
 import sys
-from unittest.mock import patch
 
 import matplotlib
 import matplotlib.pyplot as plt
@@ -184,7 +183,7 @@ def test_plot_results_components_data_structure(datadir):
         )
 
 
-def test_plot_results_wrong_axis_labels(datadir):
+def test_plot_results_wrong_axis_labels(datadir, mocker):
     """
     If the returned labels are different from the expected, then the hardcoded
     values in label_part will be wrong, causing `next` to fail on the iterator
@@ -195,12 +194,12 @@ def test_plot_results_wrong_axis_labels(datadir):
     fig = Figure()
     ax = fig.subplots()
 
-    with patch(
-        "matplotlib.axes._axes.Axes.get_legend_handles_labels"
-    ) as mock_get_legend_handles_labels:
-        mock_get_legend_handles_labels.return_value = None, ["fail"]
+    get_legend_handles_labels = mocker.patch(
+        "matplotlib.axes._axes.Axes.get_legend_handles_labels",
+        return_value=([None], ["fail"]),
+    )
 
-        with pytest.raises(StopIteration):
-            brazil.plot_results(data["testmus"], data["results"], test_size=0.05, ax=ax)
+    with pytest.raises(StopIteration):
+        brazil.plot_results(data["testmus"], data["results"], test_size=0.05, ax=ax)
 
-        mock_get_legend_handles_labels.assert_called()
+    assert get_legend_handles_labels.called