Restore input_select and test helper proposal (#6148)

* Restore input_select and test helper proposal

* DB still active

homeassistant/components/input_select.py

@@ -13,6 +13,7 @@ from homeassistant.const import ATTR_ENTITY_ID, CONF_ICON, CONF_NAME
 import homeassistant.helpers.config_validation as cv
 from homeassistant.helpers.entity import Entity
 from homeassistant.helpers.entity_component import EntityComponent
+from homeassistant.helpers.restore_state import async_get_last_state
 
 
 DOMAIN = 'input_select'
@@ -194,6 +195,16 @@ class InputSelect(Entity):
         self._options = options
         self._icon = icon
 
+    @asyncio.coroutine
+    def async_added_to_hass(self):
+        """Called when entity about to be added to hass."""
+        state = yield from async_get_last_state(self.hass, self.entity_id)
+        if not state:
+            return
+        if state.state not in self._options:
+            return
+        self._current_option = state.state
+
     @property
     def should_poll(self):
         """If entity should be polled."""

homeassistant/components/recorder/__init__.py

@@ -297,6 +297,9 @@ class Recorder(threading.Thread):
         """Tell the recorder to shut down."""
         global _INSTANCE  # pylint: disable=global-statement
         self.queue.put(None)
+        if not self.start_recording.is_set():
+            _LOGGER.warning("Recorder never started correctly")
+            self.start_recording.set()
         self.join()
         _INSTANCE = None
 

tests/common.py

@@ -15,6 +15,7 @@ from homeassistant import core as ha, loader
 from homeassistant.bootstrap import (
     setup_component, async_prepare_setup_component)
 from homeassistant.helpers.entity import ToggleEntity
+from homeassistant.helpers.restore_state import DATA_RESTORE_CACHE
 from homeassistant.util.unit_system import METRIC_SYSTEM
 import homeassistant.util.dt as date_util
 import homeassistant.util.yaml as yaml
@@ -88,6 +89,7 @@ def async_test_home_assistant(loop):
     hass = ha.HomeAssistant(loop)
 
     def async_add_job(target, *args):
+        """Add a magic mock."""
         if isinstance(target, MagicMock):
             return
         hass._async_add_job_tracking(target, *args)
@@ -459,24 +461,19 @@ def init_recorder_component(hass, add_config=None, db_ready_callback=None):
     config = dict(add_config) if add_config else {}
     config[recorder.CONF_DB_URL] = 'sqlite://'  # In memory DB
 
-    saved_recorder = recorder.Recorder
-
-    class Recorder2(saved_recorder):
-        """Recorder with a callback after db_ready."""
-
-        def _setup_connection(self):
-            """Setup the connection and run the callback."""
-            super(Recorder2, self)._setup_connection()
-            if db_ready_callback:
-                _LOGGER.debug('db_ready_callback start (db_ready not set,'
-                              'never use get_instance in the callback)')
-                db_ready_callback()
-                _LOGGER.debug('db_ready_callback completed')
-
-    with patch('homeassistant.components.recorder.Recorder',
-               side_effect=Recorder2):
-        assert setup_component(hass, recorder.DOMAIN,
-                               {recorder.DOMAIN: config})
+    assert setup_component(hass, recorder.DOMAIN,
+                           {recorder.DOMAIN: config})
     assert recorder.DOMAIN in hass.config.components
     recorder.get_instance().block_till_db_ready()
     _LOGGER.info("In-memory recorder successfully started")
+
+
+def mock_restore_cache(hass, states):
+    """Mock the DATA_RESTORE_CACHE."""
+    hass.data[DATA_RESTORE_CACHE] = {
+        state.entity_id: state for state in states}
+    _LOGGER.debug('Restore cache: %s', hass.data[DATA_RESTORE_CACHE])
+    assert len(hass.data[DATA_RESTORE_CACHE]) == len(states), \
+        "Duplicate entity_id? {}".format(states)
+    hass.state = ha.CoreState.starting
+    hass.config.components.add(recorder.DOMAIN)

tests/components/test_input_select.py

@@ -1,10 +1,12 @@
 """The tests for the Input select component."""
 # pylint: disable=protected-access
+import asyncio
 import unittest
 
-from tests.common import get_test_home_assistant
+from tests.common import get_test_home_assistant, mock_restore_cache
 
-from homeassistant.bootstrap import setup_component
+from homeassistant.core import State
+from homeassistant.bootstrap import setup_component, async_setup_component
 from homeassistant.components.input_select import (
     ATTR_OPTIONS, DOMAIN, SERVICE_SET_OPTIONS,
     select_option, select_next, select_previous)
@@ -211,3 +213,35 @@ class TestInputSelect(unittest.TestCase):
         self.hass.block_till_done()
         state = self.hass.states.get(entity_id)
         self.assertEqual('test2', state.state)
+
+
+@asyncio.coroutine
+def test_restore_state(hass):
+    """Ensure states are restored on startup."""
+    mock_restore_cache(hass, (
+        State('input_select.s1', 'last option'),
+        State('input_select.s2', 'bad option'),
+    ))
+
+    options = {
+        'options': [
+            'first option',
+            'middle option',
+            'last option',
+        ],
+        'initial': 'middle option',
+    }
+
+    yield from async_setup_component(hass, DOMAIN, {
+        DOMAIN: {
+            's1': options,
+            's2': options,
+        }})
+
+    state = hass.states.get('input_select.s1')
+    assert state
+    assert state.state == 'last option'
+
+    state = hass.states.get('input_select.s2')
+    assert state
+    assert state.state == 'middle option'

tests/helpers/test_restore_state.py

@@ -1,14 +1,18 @@
 """The tests for the Restore component."""
 import asyncio
+from datetime import timedelta
 from unittest.mock import patch, MagicMock
 
+from homeassistant.bootstrap import setup_component
 from homeassistant.const import EVENT_HOMEASSISTANT_START
-from homeassistant.core import CoreState, State
+from homeassistant.core import CoreState, split_entity_id, State
 import homeassistant.util.dt as dt_util
-
+from homeassistant.components import input_boolean, recorder
 from homeassistant.helpers.restore_state import (
     async_get_last_state, DATA_RESTORE_CACHE)
 
+from tests.common import get_test_home_assistant, init_recorder_component
+
 
 @asyncio.coroutine
 def test_caching_data(hass):
@@ -40,3 +44,67 @@ def test_caching_data(hass):
     yield from hass.async_block_till_done()
 
     assert DATA_RESTORE_CACHE not in hass.data
+
+
+def _add_data_in_last_run(entities):
+    """Add test data in the last recorder_run."""
+    # pylint: disable=protected-access
+    t_now = dt_util.utcnow() - timedelta(minutes=10)
+    t_min_1 = t_now - timedelta(minutes=20)
+    t_min_2 = t_now - timedelta(minutes=30)
+
+    recorder_runs = recorder.get_model('RecorderRuns')
+    states = recorder.get_model('States')
+    with recorder.session_scope() as session:
+        run = recorder_runs(
+            start=t_min_2,
+            end=t_now,
+            created=t_min_2
+        )
+        recorder._INSTANCE._commit(session, run)
+
+        for entity_id, state in entities.items():
+            dbstate = states(
+                entity_id=entity_id,
+                domain=split_entity_id(entity_id)[0],
+                state=state,
+                attributes='{}',
+                last_changed=t_min_1,
+                last_updated=t_min_1,
+                created=t_min_1)
+            recorder._INSTANCE._commit(session, dbstate)
+
+
+def test_filling_the_cache():
+    """Test filling the cache from the DB."""
+    test_entity_id1 = 'input_boolean.b1'
+    test_entity_id2 = 'input_boolean.b2'
+
+    hass = get_test_home_assistant()
+    hass.state = CoreState.starting
+
+    init_recorder_component(hass)
+
+    _add_data_in_last_run({
+        test_entity_id1: 'on',
+        test_entity_id2: 'off',
+    })
+
+    hass.block_till_done()
+    setup_component(hass, input_boolean.DOMAIN, {
+        input_boolean.DOMAIN: {
+            'b1': None,
+            'b2': None,
+        }})
+
+    hass.start()
+
+    state = hass.states.get('input_boolean.b1')
+    assert state
+    assert state.state == 'on'
+
+    state = hass.states.get('input_boolean.b2')
+    assert state
+    assert state.state == 'off'
+
+    hass.stop()