Add unit tests for draw features layer removal

gee-community · Jun 14, 2024 · 221e410 · 221e410
1 parent 8a93229
commit 221e410
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 7 deletions.
diff --git a/geemap/core.py b/geemap/core.py
@@ -231,7 +231,11 @@ def _handle_geometry_deleted(self, geo_json):
         if index >= 0:
             del self.geometries[index]
             del self.properties[index]
-            self._redraw_layer()
+            if self.count:
+                self._redraw_layer()
+            elif _DRAWN_FEATURES_LAYER in self.host_map.ee_layers:
+                # Remove drawn features layer if there are no geometries.
+                self.host_map.remove_layer(_DRAWN_FEATURES_LAYER)
             self._geometry_delete_dispatcher(self, geometry=geometry)
 
 
@@ -266,10 +270,6 @@ def handle_draw(_, action, geo_json):
                     self._handle_geometry_edited(geo_json)
                 elif action == "deleted":
                     self._handle_geometry_deleted(geo_json)
-                    # Remove drawn features layer if there are no geometries.
-                    if not self.count:
-                        if _DRAWN_FEATURES_LAYER in self.host_map.ee_layers:
-                            self.host_map.remove_layer(_DRAWN_FEATURES_LAYER)
             except Exception as e:
                 self.reset(clear_draw_control=False)
                 print("There was an error creating Earth Engine Feature.")

diff --git a/tests/fake_map.py b/tests/fake_map.py
@@ -84,7 +84,10 @@ def add(self, obj):
         pass
 
     def remove_layer(self, layer):
+        if isinstance(layer, str):
+            layer = self.ee_layers[layer]["ee_layer"]
         self.layers.remove(layer)
+        del self.ee_layers[layer.name]
 
     def get_layer_names(self):
         return [layer.name for layer in self.layers]

diff --git a/tests/test_core.py b/tests/test_core.py
@@ -202,8 +202,8 @@ class TestAbstractDrawControl(unittest.TestCase):
     }
 
     def setUp(self):
-        map = fake_map.FakeMap()
-        self._draw_control = TestAbstractDrawControl.TestDrawControl(map)
+        self.map = fake_map.FakeMap()
+        self._draw_control = TestAbstractDrawControl.TestDrawControl(self.map)
 
     def tearDown(self):
         pass
@@ -220,19 +220,23 @@ def test_initialization(self):
         self.assertEquals(self._draw_control.collection, fake_ee.FeatureCollection([]))
         self.assertIsNone(self._draw_control.last_feature)
         self.assertEquals(self._draw_control.count, 0)
+        self.assertFalse("Drawn Features" in self.map.ee_layers)
 
     def test_handles_creation(self):
         self._draw_control.create(self.geo_json)
         self.assertEquals(
             self._draw_control.geometries,
             [fake_ee.Geometry(self.geo_json["geometry"])],
         )
+        self.assertTrue("Drawn Features" in self.map.ee_layers)
 
     def test_handles_deletion(self):
         self._draw_control.create(self.geo_json)
+        self.assertTrue("Drawn Features" in self.map.ee_layers)
         self.assertEquals(len(self._draw_control.geometries), 1)
         self._draw_control.delete(0)
         self.assertEquals(len(self._draw_control.geometries), 0)
+        self.assertFalse("Drawn Features" in self.map.ee_layers)
 
     def test_handles_edit(self):
         self._draw_control.create(self.geo_json)
@@ -292,6 +296,7 @@ def test_reset(self):
         self._draw_control.reset(clear_draw_control=True)
         self.assertEquals(len(self._draw_control.geometries), 0)
         self.assertEquals(len(self._draw_control.geo_jsons), 0)
+        self.assertFalse("Drawn Features" in self.map.ee_layers)
 
         self._draw_control.create(self.geo_json)
         self.assertEquals(len(self._draw_control.geometries), 1)
@@ -300,6 +305,7 @@ def test_reset(self):
         self._draw_control.reset(clear_draw_control=False)
         self.assertEquals(len(self._draw_control.geometries), 0)
         self.assertEquals(len(self._draw_control.geo_jsons), 1)
+        self.assertFalse("Drawn Features" in self.map.ee_layers)
 
     def test_remove_geometry(self):
         self._draw_control.create(self.geo_json)