fix remove and re-add scenario (#25370)

This commit is contained in:
David F. Mulcahey 2019-07-21 13:00:27 -04:00 committed by Paulus Schoutsen
parent d64f1e767c
commit 7a8130cd2b
2 changed files with 12 additions and 2 deletions

View File

@ -190,6 +190,13 @@ class ZHAGateway:
if entity_id == entity_reference.reference_id:
return entity_reference
def remove_entity_reference(self, entity):
"""Remove entity reference for given entity_id if found."""
if entity.zha_device.ieee in self.device_registry:
entity_refs = self.device_registry.get(entity.zha_device.ieee)
self.device_registry[entity.zha_device.ieee] = [
e for e in entity_refs if e.reference_id != entity.entity_id]
@property
def devices(self):
"""Return devices."""

View File

@ -50,7 +50,7 @@ class ZhaEntity(RestoreEntity, entity.Entity):
self._available = False
self._component = kwargs['component']
self._unsubs = []
self.remove_future = asyncio.Future()
self.remove_future = None
for channel in channels:
self.cluster_channels[channel.name] = channel
@ -123,6 +123,7 @@ class ZhaEntity(RestoreEntity, entity.Entity):
async def async_added_to_hass(self):
"""Run when about to be added to hass."""
await super().async_added_to_hass()
self.remove_future = asyncio.Future()
await self.async_check_recently_seen()
await self.async_accept_signal(
None, "{}_{}".format(self.zha_device.available_signal, 'entity'),
@ -151,8 +152,10 @@ class ZhaEntity(RestoreEntity, entity.Entity):
async def async_will_remove_from_hass(self) -> None:
"""Disconnect entity object when removed."""
for unsub in self._unsubs:
for unsub in self._unsubs[:]:
unsub()
self._unsubs.remove(unsub)
self.zha_device.gateway.remove_entity_reference(self)
self.remove_future.set_result(True)
@callback