Restore manual alarm-control-panel state using async_get_last_state

homeassistant/components/alarm_control_panel/manual.py

@@ -21,6 +21,7 @@
 import homeassistant.helpers.config_validation as cv
 from homeassistant.helpers.event import track_point_in_time
 import homeassistant.util.dt as dt_util
+from homeassistant.helpers.restore_state import async_get_last_state
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -306,3 +307,10 @@ def device_state_attributes(self):
             state_attr[ATTR_POST_PENDING_STATE] = self._state
 
         return state_attr
+
+    async def async_added_to_hass(self):
+        """Run when entity about to be added to hass."""
+        state = await async_get_last_state(self.hass, self.entity_id)
+        if state:
+            self._state = state.state
+            self._state_ts = state.last_updated

tests/components/alarm_control_panel/test_manual.py

@@ -3,18 +3,18 @@
 import unittest
 from unittest.mock import patch, MagicMock
 from homeassistant.components.alarm_control_panel import demo
-
-
-from homeassistant.setup import setup_component
+import asyncio
+from homeassistant.setup import setup_component, async_setup_component
 from homeassistant.const import (
     STATE_ALARM_DISARMED, STATE_ALARM_ARMED_HOME, STATE_ALARM_ARMED_AWAY,
     STATE_ALARM_ARMED_NIGHT, STATE_ALARM_ARMED_CUSTOM_BYPASS,
     STATE_ALARM_PENDING, STATE_ALARM_TRIGGERED)
 from homeassistant.components import alarm_control_panel
 import homeassistant.util.dt as dt_util
-
-from tests.common import fire_time_changed, get_test_home_assistant
+from tests.common import (fire_time_changed, get_test_home_assistant, 
+        mock_component, mock_restore_cache)
 from tests.components.alarm_control_panel import common
+from homeassistant.core import State, CoreState
 
 CODE = 'HELLO_CODE'
 
@@ -1319,3 +1319,50 @@ def test_arm_away_after_disabled_disarmed(self):
 
         state = self.hass.states.get(entity_id)
         self.assertEqual(STATE_ALARM_TRIGGERED, state.state)
+
+@asyncio.coroutine
+def test_restore_armed_state(hass):
+    """Ensure armed state is restored on startup."""
+    mock_restore_cache(hass, (
+        State('alarm_control_panel.test', STATE_ALARM_ARMED_AWAY),
+        ))
+
+    hass.state = CoreState.starting
+    mock_component(hass, 'recorder')
+
+    yield from async_setup_component(hass, alarm_control_panel.DOMAIN,
+        {'alarm_control_panel': {
+            'platform': 'manual',
+            'name': 'test',
+            'pending_time': 0,
+            'trigger_time': 0,
+            'disarm_after_trigger': False
+        }})
+
+    state = hass.states.get('alarm_control_panel.test')
+    assert state
+    assert state.state == STATE_ALARM_ARMED_AWAY
+
+@asyncio.coroutine
+def test_restore_disarmed_state(hass):
+    """Ensure disarmed state is restored on startup."""
+    mock_restore_cache(hass, (
+        State('alarm_control_panel.test', STATE_ALARM_DISARMED),
+        ))
+
+    hass.state = CoreState.starting
+    mock_component(hass, 'recorder')
+
+    yield from async_setup_component(hass, alarm_control_panel.DOMAIN,
+        {'alarm_control_panel': {
+            'platform': 'manual',
+            'name': 'test',
+            'pending_time': 0,
+            'trigger_time': 0,
+            'disarm_after_trigger': False
+        }})
+
+    state = hass.states.get('alarm_control_panel.test')
+    assert state
+    assert state.state == STATE_ALARM_DISARMED
+