Restore states when removing/adding entities (#18890)
This commit is contained in:
parent
bbb40fde84
commit
a10cbadb57
2 changed files with 30 additions and 4 deletions
|
@ -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."""
|
||||
|
|
|
@ -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'
|
||||
|
|
Loading…
Add table
Add a link
Reference in a new issue