Restore states when removing/adding entities (#18890)

This commit is contained in:
Adam Mills 2018-12-02 04:51:15 -05:00 committed by Paulus Schoutsen
parent bbb40fde84
commit a10cbadb57
2 changed files with 30 additions and 4 deletions

View file

@ -165,13 +165,19 @@ class RestoreStateData():
self.async_dump_states()))
@callback
def async_register_entity(self, entity_id: str) -> None:
def async_restore_entity_added(self, entity_id: str) -> None:
"""Store this entity's state when hass is shutdown."""
self.entity_ids.add(entity_id)
@callback
def async_unregister_entity(self, entity_id: str) -> None:
def async_restore_entity_removed(self, entity_id: str) -> None:
"""Unregister this entity from saving state."""
# When an entity is being removed from hass, store its last state. This
# allows us to support state restoration if the entity is removed, then
# re-added while hass is still running.
self.last_states[entity_id] = StoredState(
self.hass.states.get(entity_id), dt_util.utcnow())
self.entity_ids.remove(entity_id)
@ -184,7 +190,7 @@ class RestoreEntity(Entity):
super().async_added_to_hass(),
RestoreStateData.async_get_instance(self.hass),
)
data.async_register_entity(self.entity_id)
data.async_restore_entity_added(self.entity_id)
async def async_will_remove_from_hass(self) -> None:
"""Run when entity will be removed from hass."""
@ -192,7 +198,7 @@ class RestoreEntity(Entity):
super().async_will_remove_from_hass(),
RestoreStateData.async_get_instance(self.hass),
)
data.async_unregister_entity(self.entity_id)
data.async_restore_entity_removed(self.entity_id)
async def async_get_last_state(self) -> Optional[State]:
"""Get the entity state from the previous run."""

View file

@ -198,3 +198,23 @@ async def test_load_error(hass):
state = await entity.async_get_last_state()
assert state is None
async def test_state_saved_on_remove(hass):
"""Test that we save entity state on removal."""
entity = RestoreEntity()
entity.hass = hass
entity.entity_id = 'input_boolean.b0'
await entity.async_added_to_hass()
hass.states.async_set('input_boolean.b0', 'on')
data = await RestoreStateData.async_get_instance(hass)
# No last states should currently be saved
assert not data.last_states
await entity.async_will_remove_from_hass()
# We should store the input boolean state when it is removed
assert data.last_states['input_boolean.b0'].state.state == 'on'