From 0d6ee255831adae452af355c025497c0f07aa296 Mon Sep 17 00:00:00 2001
From: Oscar Gustafsson <oscar.gustafsson@gmail.com>
Date: Sat, 15 Oct 2022 17:46:54 +0200
Subject: [PATCH] Add tests for ToolManager

---
 lib/matplotlib/backend_managers.py         |  5 +-
 lib/matplotlib/tests/test_backend_bases.py | 54 ++++++++++++++++++++--
 2 files changed, 52 insertions(+), 7 deletions(-)

diff --git a/lib/matplotlib/backend_managers.py b/lib/matplotlib/backend_managers.py
index 7c3f7b92106f..e7dd748325e5 100644
--- a/lib/matplotlib/backend_managers.py
+++ b/lib/matplotlib/backend_managers.py
@@ -185,7 +185,7 @@ def update_keymap(self, name, key):
             Keys to associate with the tool.
         """
         if name not in self._tools:
-            raise KeyError(f'{name} not in Tools')
+            raise KeyError(f'{name!r} not in Tools')
         self._remove_keys(name)
         if isinstance(key, str):
             key = [key]
@@ -404,6 +404,7 @@ def get_tool(self, name, warn=True):
             return name
         if name not in self._tools:
             if warn:
-                _api.warn_external(f"ToolManager does not control tool {name}")
+                _api.warn_external(
+                    f"ToolManager does not control tool {name!r}")
             return None
         return self._tools[name]
diff --git a/lib/matplotlib/tests/test_backend_bases.py b/lib/matplotlib/tests/test_backend_bases.py
index 231a3e044705..96d2f3afcd55 100644
--- a/lib/matplotlib/tests/test_backend_bases.py
+++ b/lib/matplotlib/tests/test_backend_bases.py
@@ -4,6 +4,7 @@
 from matplotlib.backend_bases import (
     FigureCanvasBase, LocationEvent, MouseButton, MouseEvent,
     NavigationToolbar2, RendererBase)
+from matplotlib.backend_tools import RubberbandBase
 from matplotlib.figure import Figure
 from matplotlib.testing._markers import needs_pgf_xelatex
 import matplotlib.pyplot as plt
@@ -12,6 +13,12 @@
 import pytest
 
 
+_EXPECTED_WARNING_TOOLMANAGER = (
+    r"Treat the new Tool classes introduced in "
+    r"v[0-9]*.[0-9]* as experimental for now; "
+    "the API and rcParam may change in future versions.")
+
+
 def test_uses_per_path():
     id = transforms.Affine2D()
     paths = [path.Path.unit_regular_polygon(i) for i in range(3, 7)]
@@ -247,11 +254,7 @@ def test_interactive_colorbar(plot_func, orientation, tool, button, expected):
 
 
 def test_toolbar_zoompan():
-    expected_warning_regex = (
-        r"Treat the new Tool classes introduced in "
-        r"v[0-9]*.[0-9]* as experimental for now; "
-        "the API and rcParam may change in future versions.")
-    with pytest.warns(UserWarning, match=expected_warning_regex):
+    with pytest.warns(UserWarning, match=_EXPECTED_WARNING_TOOLMANAGER):
         plt.rcParams['toolbar'] = 'toolmanager'
     ax = plt.gca()
     assert ax.get_navigate_mode() is None
@@ -349,3 +352,44 @@ def test_interactive_pan(key, mouseend, expectedxlim, expectedylim):
     # Should be close, but won't be exact due to screen integer resolution
     assert tuple(ax.get_xlim()) == pytest.approx(expectedxlim, abs=0.02)
     assert tuple(ax.get_ylim()) == pytest.approx(expectedylim, abs=0.02)
+
+
+def test_toolmanager_remove():
+    with pytest.warns(UserWarning, match=_EXPECTED_WARNING_TOOLMANAGER):
+        plt.rcParams['toolbar'] = 'toolmanager'
+    fig = plt.gcf()
+    initial_len = len(fig.canvas.manager.toolmanager.tools)
+    assert 'forward' in fig.canvas.manager.toolmanager.tools
+    fig.canvas.manager.toolmanager.remove_tool('forward')
+    assert len(fig.canvas.manager.toolmanager.tools) == initial_len - 1
+    assert 'forward' not in fig.canvas.manager.toolmanager.tools
+
+
+def test_toolmanager_get_tool():
+    with pytest.warns(UserWarning, match=_EXPECTED_WARNING_TOOLMANAGER):
+        plt.rcParams['toolbar'] = 'toolmanager'
+    fig = plt.gcf()
+    rubberband = fig.canvas.manager.toolmanager.get_tool('rubberband')
+    assert isinstance(rubberband, RubberbandBase)
+    assert fig.canvas.manager.toolmanager.get_tool(rubberband) is rubberband
+    with pytest.warns(UserWarning,
+                      match="ToolManager does not control tool 'foo'"):
+        assert fig.canvas.manager.toolmanager.get_tool('foo') is None
+    assert fig.canvas.manager.toolmanager.get_tool('foo', warn=False) is None
+
+    with pytest.warns(UserWarning,
+                      match="ToolManager does not control tool 'foo'"):
+        assert fig.canvas.manager.toolmanager.trigger_tool('foo') is None
+
+
+def test_toolmanager_update_keymap():
+    with pytest.warns(UserWarning, match=_EXPECTED_WARNING_TOOLMANAGER):
+        plt.rcParams['toolbar'] = 'toolmanager'
+    fig = plt.gcf()
+    assert 'v' in fig.canvas.manager.toolmanager.get_tool_keymap('forward')
+    with pytest.warns(UserWarning,
+                      match="Key c changed from back to forward"):
+        fig.canvas.manager.toolmanager.update_keymap('forward', 'c')
+    assert fig.canvas.manager.toolmanager.get_tool_keymap('forward') == ['c']
+    with pytest.raises(KeyError, match="'foo' not in Tools"):
+        fig.canvas.manager.toolmanager.update_keymap('foo', 'c')