mirror of
https://github.com/home-assistant/core
synced 2024-09-06 10:29:55 +02:00
Correctly store removed entities for restore state (#25073)
* Correctly store removed entities for restore state * Lint * Do not assume about set encoding
This commit is contained in:
parent
312fceeaf6
commit
073327831f
@ -177,12 +177,43 @@ class RestoreStateData():
|
||||
# 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())
|
||||
state = self.hass.states.get(entity_id)
|
||||
# To fully mimic all the attribute data types when loaded from storage,
|
||||
# we're going to serialize it to JSON and then re-load it.
|
||||
if state is not None:
|
||||
state = State.from_dict(_encode_complex(state.as_dict()))
|
||||
|
||||
self.last_states[entity_id] = StoredState(state, dt_util.utcnow())
|
||||
|
||||
self.entity_ids.remove(entity_id)
|
||||
|
||||
|
||||
def _encode(value):
|
||||
"""Little helper to JSON encode a value."""
|
||||
try:
|
||||
return JSONEncoder.default(None, value)
|
||||
except TypeError:
|
||||
return value
|
||||
|
||||
|
||||
def _encode_complex(value):
|
||||
"""Recursively encode all values with the JSONEncoder."""
|
||||
if isinstance(value, dict):
|
||||
return {
|
||||
_encode(key): _encode_complex(value)
|
||||
for key, value in value.items()
|
||||
}
|
||||
elif isinstance(value, list):
|
||||
return [_encode_complex(val) for val in value]
|
||||
|
||||
new_value = _encode(value)
|
||||
|
||||
if isinstance(new_value, type(value)):
|
||||
return new_value
|
||||
|
||||
return _encode_complex(new_value)
|
||||
|
||||
|
||||
class RestoreEntity(Entity):
|
||||
"""Mixin class for restoring previous entity state."""
|
||||
|
||||
|
@ -1,6 +1,8 @@
|
||||
"""The tests for the Restore component."""
|
||||
from datetime import datetime
|
||||
|
||||
from asynctest import patch
|
||||
|
||||
from homeassistant.const import EVENT_HOMEASSISTANT_START
|
||||
from homeassistant.core import CoreState, State
|
||||
from homeassistant.exceptions import HomeAssistantError
|
||||
@ -10,7 +12,6 @@ from homeassistant.helpers.restore_state import (
|
||||
STORAGE_KEY)
|
||||
from homeassistant.util import dt as dt_util
|
||||
|
||||
from asynctest import patch
|
||||
|
||||
from tests.common import mock_coro
|
||||
|
||||
@ -208,7 +209,12 @@ async def test_state_saved_on_remove(hass):
|
||||
entity.entity_id = 'input_boolean.b0'
|
||||
await entity.async_internal_added_to_hass()
|
||||
|
||||
hass.states.async_set('input_boolean.b0', 'on')
|
||||
now = dt_util.utcnow()
|
||||
hass.states.async_set('input_boolean.b0', 'on', {
|
||||
'complicated': {
|
||||
'value': {1, 2, now}
|
||||
}
|
||||
})
|
||||
|
||||
data = await RestoreStateData.async_get_instance(hass)
|
||||
|
||||
@ -218,7 +224,12 @@ async def test_state_saved_on_remove(hass):
|
||||
await entity.async_remove()
|
||||
|
||||
# We should store the input boolean state when it is removed
|
||||
assert data.last_states['input_boolean.b0'].state.state == 'on'
|
||||
state = data.last_states['input_boolean.b0'].state
|
||||
assert state.state == 'on'
|
||||
assert isinstance(state.attributes['complicated']['value'], list)
|
||||
assert set(state.attributes['complicated']['value']) == {
|
||||
1, 2, now.isoformat()
|
||||
}
|
||||
|
||||
|
||||
async def test_restoring_invalid_entity_id(hass, hass_storage):
|
||||
|
Loading…
Reference in New Issue
Block a user