Restore states through a JSON store instead of recorder (#17270)

* Restore states through a JSON store

* Accept entity_id directly in restore state helper

* Keep states stored between runs for a limited time

* Remove warning

tests/common.py

@@ -114,8 +114,7 @@ def get_test_home_assistant():
 
 
 # pylint: disable=protected-access
-@asyncio.coroutine
-def async_test_home_assistant(loop):
+async def async_test_home_assistant(loop):
     """Return a Home Assistant object pointing at test config dir."""
     hass = ha.HomeAssistant(loop)
     hass.config.async_load = Mock()
@@ -168,13 +167,12 @@ def async_test_home_assistant(loop):
     # Mock async_start
     orig_start = hass.async_start
 
-    @asyncio.coroutine
-    def mock_async_start():
+    async def mock_async_start():
         """Start the mocking."""
         # We only mock time during tests and we want to track tasks
         with patch('homeassistant.core._async_create_timer'), \
                 patch.object(hass, 'async_stop_track_tasks'):
-            yield from orig_start()
+            await orig_start()
 
     hass.async_start = mock_async_start
 
@@ -715,14 +713,20 @@ def init_recorder_component(hass, add_config=None):
 
 def mock_restore_cache(hass, states):
     """Mock the DATA_RESTORE_CACHE."""
-    key = restore_state.DATA_RESTORE_CACHE
-    hass.data[key] = {
+    key = restore_state.DATA_RESTORE_STATE_TASK
+    data = restore_state.RestoreStateData(hass)
+
+    data.last_states = {
         state.entity_id: state for state in states}
-    _LOGGER.debug('Restore cache: %s', hass.data[key])
-    assert len(hass.data[key]) == len(states), \
+    _LOGGER.debug('Restore cache: %s', data.last_states)
+    assert len(data.last_states) == len(states), \
         "Duplicate entity_id? {}".format(states)
-    hass.state = ha.CoreState.starting
-    mock_component(hass, recorder.DOMAIN)
+
+    async def get_restore_state_data() -> restore_state.RestoreStateData:
+        return data
+
+    # Patch the singleton task in hass.data to return our new RestoreStateData
+    hass.data[key] = hass.async_create_task(get_restore_state_data())
 
 
 class MockDependency:
@@ -846,9 +850,10 @@ def mock_storage(data=None):
 
     def mock_write_data(store, path, data_to_write):
         """Mock version of write data."""
-        # To ensure that the data can be serialized
         _LOGGER.info('Writing data to %s: %s', store.key, data_to_write)
-        data[store.key] = json.loads(json.dumps(data_to_write))
+        # To ensure that the data can be serialized
+        data[store.key] = json.loads(json.dumps(
+            data_to_write, cls=store._encoder))
 
     with patch('homeassistant.helpers.storage.Store._async_load',
                side_effect=mock_async_load, autospec=True), \