Fix automation failing to restore state (#24390)

* Fix automation off

* Fix tests

homeassistant/components/automation/__init__.py

@@ -190,6 +190,7 @@ def __init__(self, automation_id, name, async_attach_triggers, cond_func,
         self._last_triggered = None
         self._hidden = hidden
         self._initial_state = initial_state
+        self._is_enabled = False
 
     @property
     def name(self):
@@ -216,7 +217,8 @@ def hidden(self) -> bool:
     @property
     def is_on(self) -> bool:
         """Return True if entity is on."""
-        return self._async_detach_triggers is not None
+        return (self._async_detach_triggers is not None or
+                self._is_enabled)
 
     async def async_added_to_hass(self) -> None:
         """Startup with initial state or previous state."""
@@ -239,37 +241,16 @@ async def async_added_to_hass(self) -> None:
                               "initial state", self.entity_id,
                               enable_automation)
 
-        if not enable_automation:
-            return
-
-        # HomeAssistant is starting up
-        if self.hass.state == CoreState.not_running:
-            async def async_enable_automation(event):
-                """Start automation on startup."""
-                await self.async_enable()
-
-            self.hass.bus.async_listen_once(
-                EVENT_HOMEASSISTANT_START, async_enable_automation)
-
-        # HomeAssistant is running
-        else:
+        if enable_automation:
             await self.async_enable()
 
     async def async_turn_on(self, **kwargs) -> None:
         """Turn the entity on and update the state."""
-        if self.is_on:
-            return
-
         await self.async_enable()
 
     async def async_turn_off(self, **kwargs) -> None:
         """Turn the entity off."""
-        if not self.is_on:
-            return
-
-        self._async_detach_triggers()
-        self._async_detach_triggers = None
-        await self.async_update_ha_state()
+        await self.async_disable()
 
     async def async_trigger(self, variables, skip_condition=False,
                             context=None):
@@ -296,19 +277,51 @@ async def async_trigger(self, variables, skip_condition=False,
     async def async_will_remove_from_hass(self):
         """Remove listeners when removing automation from HASS."""
         await super().async_will_remove_from_hass()
-        await self.async_turn_off()
+        await self.async_disable()
 
     async def async_enable(self):
         """Enable this automation entity.
 
         This method is a coroutine.
         """
-        if self.is_on:
+        if self._is_enabled:
             return
 
-        self._async_detach_triggers = await self._async_attach_triggers(
-            self.async_trigger)
-        await self.async_update_ha_state()
+        self._is_enabled = True
+
+        # HomeAssistant is starting up
+        if self.hass.state != CoreState.not_running:
+            self._async_detach_triggers = await self._async_attach_triggers(
+                self.async_trigger)
+            self.async_write_ha_state()
+            return
+
+        async def async_enable_automation(event):
+            """Start automation on startup."""
+            # Don't do anything if no longer enabled or already attached
+            if (not self._is_enabled or
+                    self._async_detach_triggers is not None):
+                return
+
+            self._async_detach_triggers = await self._async_attach_triggers(
+                self.async_trigger)
+
+        self.hass.bus.async_listen_once(
+            EVENT_HOMEASSISTANT_START, async_enable_automation)
+        self.async_write_ha_state()
+
+    async def async_disable(self):
+        """Disable the automation entity."""
+        if not self._is_enabled:
+            return
+
+        self._is_enabled = False
+
+        if self._async_detach_triggers is not None:
+            self._async_detach_triggers()
+            self._async_detach_triggers = None
+
+        self.async_write_ha_state()
 
     @property
     def device_state_attributes(self):

tests/components/automation/test_homeassistant.py

@@ -29,7 +29,7 @@ def test_if_fires_on_hass_start(hass):
 
     res = yield from async_setup_component(hass, automation.DOMAIN, config)
     assert res
-    assert not automation.is_on(hass, 'automation.hello')
+    assert automation.is_on(hass, 'automation.hello')
     assert len(calls) == 0
 
     yield from hass.async_start()
@@ -64,7 +64,7 @@ def test_if_fires_on_hass_shutdown(hass):
         }
     })
     assert res
-    assert not automation.is_on(hass, 'automation.hello')
+    assert automation.is_on(hass, 'automation.hello')
     assert len(calls) == 0
 
     yield from hass.async_start()

tests/components/automation/test_init.py

@@ -696,12 +696,12 @@ def test_initial_value_off(hass):
     assert len(calls) == 0
 
 
-@asyncio.coroutine
-def test_initial_value_on(hass):
+async def test_initial_value_on(hass):
     """Test initial value on."""
+    hass.state = CoreState.not_running
     calls = async_mock_service(hass, 'test', 'automation')
 
-    res = yield from async_setup_component(hass, automation.DOMAIN, {
+    assert await async_setup_component(hass, automation.DOMAIN, {
         automation.DOMAIN: {
             'alias': 'hello',
             'initial_state': 'on',
@@ -715,23 +715,23 @@ def test_initial_value_on(hass):
             }
         }
     })
-    assert res
     assert automation.is_on(hass, 'automation.hello')
 
+    await hass.async_start()
     hass.bus.async_fire('test_event')
-    yield from hass.async_block_till_done()
+    await hass.async_block_till_done()
     assert len(calls) == 1
 
 
-@asyncio.coroutine
-def test_initial_value_off_but_restore_on(hass):
+async def test_initial_value_off_but_restore_on(hass):
     """Test initial value off and restored state is turned on."""
+    hass.state = CoreState.not_running
     calls = async_mock_service(hass, 'test', 'automation')
     mock_restore_cache(hass, (
         State('automation.hello', STATE_ON),
     ))
 
-    res = yield from async_setup_component(hass, automation.DOMAIN, {
+    await async_setup_component(hass, automation.DOMAIN, {
         automation.DOMAIN: {
             'alias': 'hello',
             'initial_state': 'off',
@@ -745,11 +745,11 @@ def test_initial_value_off_but_restore_on(hass):
             }
         }
     })
-    assert res
     assert not automation.is_on(hass, 'automation.hello')
 
+    await hass.async_start()
     hass.bus.async_fire('test_event')
-    yield from hass.async_block_till_done()
+    await hass.async_block_till_done()
     assert len(calls) == 0
 
 
@@ -858,7 +858,7 @@ def test_automation_not_trigger_on_bootstrap(hass):
         }
     })
     assert res
-    assert not automation.is_on(hass, 'automation.hello')
+    assert automation.is_on(hass, 'automation.hello')
 
     hass.bus.async_fire('test_event')
     yield from hass.async_block_till_done()