From 75ffea1dbb24c250193e78c8551e78d89da3b870 Mon Sep 17 00:00:00 2001
From: Sufyan Abbasi <sufy@google.com>
Date: Fri, 6 Oct 2023 06:59:11 +0000
Subject: [PATCH] Add unit tests for toolbar cleanup functionality

---
 tests/test_toolbar.py | 50 ++++++++++++++++++++++++++++++++++++++++---
 1 file changed, 47 insertions(+), 3 deletions(-)

diff --git a/tests/test_toolbar.py b/tests/test_toolbar.py
index 16a5822b4a..aed699e1d7 100644
--- a/tests/test_toolbar.py
+++ b/tests/test_toolbar.py
@@ -2,14 +2,15 @@
 
 """Tests for `map_widgets` module."""
 
-
 import unittest
+
+from dataclasses import dataclass
 from unittest.mock import patch, Mock
 
 import ipywidgets
 
 import geemap
-from geemap.toolbar import Toolbar
+from geemap.toolbar import Toolbar, _cleanup_toolbar_item
 from tests import fake_map, utils
 
 
@@ -32,6 +33,7 @@ def _query_tool_grid(self, toolbar):
     def setUp(self) -> None:
         self.callback_calls = 0
         self.last_called_with_selected = None
+        self.last_called_item = None
         self.item = Toolbar.Item(
             icon="info", tooltip="dummy item", callback=self.dummy_callback
         )
@@ -47,9 +49,10 @@ def tearDown(self) -> None:
         patch.stopall()
         return super().tearDown()
 
-    def dummy_callback(self, m, selected):
+    def dummy_callback(self, m, selected, item):
         del m
         self.last_called_with_selected = selected
+        self.last_called_item = item
         self.callback_calls += 1
 
     def test_no_tools_throws(self):
@@ -113,18 +116,21 @@ def test_triggers_callbacks(self):
         map = geemap.Map(ee_initialize=False)
         toolbar = Toolbar(map, [self.item, self.no_reset_item])
         self.assertIsNone(self.last_called_with_selected)
+        self.assertIsNone(self.last_called_item)
 
         # Select first tool, which resets.
         toolbar.all_widgets[0].value = True
         self.assertFalse(self.last_called_with_selected)  # was reset by callback
         self.assertEqual(self.callback_calls, 2)
         self.assertFalse(toolbar.all_widgets[0].value)
+        self.assertEqual(self.item, self.last_called_item)
 
         # Select second tool, which does not reset.
         toolbar.all_widgets[1].value = True
         self.assertTrue(self.last_called_with_selected)
         self.assertEqual(self.callback_calls, 3)
         self.assertTrue(toolbar.all_widgets[1].value)
+        self.assertEqual(self.no_reset_item, self.last_called_item)
 
     def test_layers_toggle_callback(self):
         """Verifies the on_layers_toggled callback is triggered."""
@@ -154,3 +160,41 @@ def test_accessory_widget(self):
                 toolbar, ipywidgets.ToggleButton, lambda c: c.tooltip == "test-button"
             )
         )
+
+    @dataclass
+    class TestWidget:
+        selected_count = 0
+        cleanup_count = 0
+
+        def cleanup(self):
+            self.cleanup_count += 1
+
+    def test_cleanup_toolbar_item_decorator(self):
+        widget = TestToolbar.TestWidget()
+
+        @_cleanup_toolbar_item
+        def callback(m, selected, item):
+            widget.selected_count += 1
+            return widget
+
+        item = Toolbar.Item(
+            icon="info", tooltip="dummy item", callback=callback, reset=False
+        )
+        map_fake = fake_map.FakeMap()
+        toolbar = Toolbar(map_fake, [item])
+        toolbar.all_widgets[0].value = True
+        self.assertEqual(1, widget.selected_count)
+        self.assertEqual(0, widget.cleanup_count)
+
+        toolbar.all_widgets[0].value = False
+        self.assertEqual(1, widget.selected_count)
+        self.assertEqual(1, widget.cleanup_count)
+
+        toolbar.all_widgets[0].value = True
+        self.assertEqual(2, widget.selected_count)
+        self.assertEqual(1, widget.cleanup_count)
+
+        widget.cleanup()
+        self.assertEqual(2, widget.selected_count)
+        self.assertEqual(3, widget.cleanup_count)
+        self.assertFalse(toolbar.all_widgets[0].value)