Allow removing entity registry items twice (#146519)

homeassistant/helpers/entity_registry.py

@@ -1045,6 +1045,14 @@ class EntityRegistry(BaseRegistry):
     def async_remove(self, entity_id: str) -> None:
         """Remove an entity from registry."""
         self.hass.verify_event_loop_thread("entity_registry.async_remove")
+        if entity_id not in self.entities:
+            # Allow attempts to remove an entity which does not exist. If this is
+            # not allowed, there will be races during cleanup where we iterate over
+            # lists of entities to remove, but there are listeners for entity
+            # registry events which delete entities at the same time.
+            # For example, if we clean up entities A and B, there might be a listener
+            # which deletes entity B when entity A is being removed.
+            return
         entity = self.entities.pop(entity_id)
         config_entry_id = entity.config_entry_id
         key = (entity.domain, entity.platform, entity.unique_id)

tests/helpers/test_entity_registry.py

@@ -289,6 +289,24 @@ def test_get_or_create_suggested_object_id_conflict_existing(
     assert entry.entity_id == "light.hue_1234_2"
 
 
+def test_remove(entity_registry: er.EntityRegistry) -> None:
+    """Test that we can remove an item."""
+    entry = entity_registry.async_get_or_create("light", "hue", "1234")
+
+    assert not entity_registry.deleted_entities
+    assert list(entity_registry.entities) == [entry.entity_id]
+
+    # Remove the item
+    entity_registry.async_remove(entry.entity_id)
+    assert list(entity_registry.deleted_entities) == [("light", "hue", "1234")]
+    assert not entity_registry.entities
+
+    # Remove the item again
+    entity_registry.async_remove(entry.entity_id)
+    assert list(entity_registry.deleted_entities) == [("light", "hue", "1234")]
+    assert not entity_registry.entities
+
+
 def test_create_triggers_save(entity_registry: er.EntityRegistry) -> None:
     """Test that registering entry triggers a save."""
     with patch.object(entity_registry, "async_schedule_save") as mock_schedule_save: