Untangle entity remove code from entity platform

home-assistant · Oct 25, 2018 · fdfb7c2 · fdfb7c2
1 parent a68e391
commit fdfb7c2
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 19 deletions.
diff --git a/homeassistant/helpers/entity.py b/homeassistant/helpers/entity.py
@@ -361,16 +361,16 @@ def async_on_remove(self, func):
             self._on_remove = []
         self._on_remove.append(func)
 
-    async def async_remove(self):
-        """Remove entity from Home Assistant."""
+    async def async_will_remove_from_hass(self):
+        """Do things when entity is removed from hass."""
         if self._on_remove is not None:
             while self._on_remove:
                 self._on_remove.pop()()
 
-        if self.platform is not None:
-            await self.platform.async_remove_entity(self.entity_id)
-        else:
-            self.hass.states.async_remove(self.entity_id)
+    async def async_remove(self):
+        """Remove entity from Home Assistant."""
+        await self.async_will_remove_from_hass()
+        self.hass.states.async_remove(self.entity_id)
 
     @callback
     def async_registry_updated(self, old, new):

diff --git a/homeassistant/helpers/entity_platform.py b/homeassistant/helpers/entity_platform.py
@@ -345,8 +345,10 @@ async def _async_add_entity(self, entity, update_before_add,
             raise HomeAssistantError(
                 msg)
 
-        self.entities[entity.entity_id] = entity
-        component_entities.add(entity.entity_id)
+        entity_id = entity.entity_id
+        self.entities[entity_id] = entity
+        component_entities.add(entity_id)
+        entity.async_on_remove(lambda: self.entities.pop(entity_id))
 
         if hasattr(entity, 'async_added_to_hass'):
             await entity.async_added_to_hass()
@@ -365,7 +367,7 @@ async def async_reset(self):
         if not self.entities:
             return
 
-        tasks = [self._async_remove_entity(entity_id)
+        tasks = [self.async_remove_entity(entity_id)
                  for entity_id in self.entities]
 
         await asyncio.wait(tasks, loop=self.hass.loop)
@@ -376,7 +378,7 @@ async def async_reset(self):
 
     async def async_remove_entity(self, entity_id):
         """Remove entity id from platform."""
-        await self._async_remove_entity(entity_id)
+        await self.entities[entity_id].async_remove()
 
         # Clean up polling job if no longer needed
         if (self._async_unsub_polling is not None and
@@ -385,15 +387,6 @@ async def async_remove_entity(self, entity_id):
             self._async_unsub_polling()
             self._async_unsub_polling = None
 
-    async def _async_remove_entity(self, entity_id):
-        """Remove entity id from platform."""
-        entity = self.entities.pop(entity_id)
-
-        if hasattr(entity, 'async_will_remove_from_hass'):
-            await entity.async_will_remove_from_hass()
-
-        self.hass.states.async_remove(entity_id)
-
     async def _update_entity_states(self, now):
         """Update the states of all the polling entities.