diff --git a/homeassistant/helpers/restore_state.py b/homeassistant/helpers/restore_state.py index 291bc6d1a0a..2f9f7991ca1 100644 --- a/homeassistant/helpers/restore_state.py +++ b/homeassistant/helpers/restore_state.py @@ -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.""" diff --git a/tests/helpers/test_restore_state.py b/tests/helpers/test_restore_state.py index ff38e8fb763..229e5b1dc1b 100644 --- a/tests/helpers/test_restore_state.py +++ b/tests/helpers/test_restore_state.py @@ -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):