Correctly store removed entities for restore state (#25073)

* Correctly store removed entities for restore state

* Lint

* Do not assume about set encoding

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."""
 

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):