Add optimistic mode to template switch (#31637)

homeassistant/components/template/switch.py

@@ -24,6 +24,7 @@ from homeassistant.exceptions import TemplateError
 import homeassistant.helpers.config_validation as cv
 from homeassistant.helpers.entity import async_generate_entity_id
 from homeassistant.helpers.event import async_track_state_change
+from homeassistant.helpers.restore_state import RestoreEntity
 from homeassistant.helpers.script import Script
 
 from . import extract_entities, initialise_templates
@@ -37,7 +38,7 @@ OFF_ACTION = "turn_off"
 
 SWITCH_SCHEMA = vol.Schema(
     {
-        vol.Required(CONF_VALUE_TEMPLATE): cv.template,
+        vol.Optional(CONF_VALUE_TEMPLATE): cv.template,
         vol.Optional(CONF_ICON_TEMPLATE): cv.template,
         vol.Optional(CONF_ENTITY_PICTURE_TEMPLATE): cv.template,
         vol.Optional(CONF_AVAILABILITY_TEMPLATE): cv.template,
@@ -59,7 +60,7 @@ async def async_setup_platform(hass, config, async_add_entities, discovery_info=
 
     for device, device_config in config[CONF_SWITCHES].items():
         friendly_name = device_config.get(ATTR_FRIENDLY_NAME, device)
-        state_template = device_config[CONF_VALUE_TEMPLATE]
+        state_template = device_config.get(CONF_VALUE_TEMPLATE)
         icon_template = device_config.get(CONF_ICON_TEMPLATE)
         entity_picture_template = device_config.get(CONF_ENTITY_PICTURE_TEMPLATE)
         availability_template = device_config.get(CONF_AVAILABILITY_TEMPLATE)
@@ -96,7 +97,7 @@ async def async_setup_platform(hass, config, async_add_entities, discovery_info=
     async_add_entities(switches)
 
 
-class SwitchTemplate(SwitchEntity):
+class SwitchTemplate(SwitchEntity, RestoreEntity):
     """Representation of a Template switch."""
 
     def __init__(
@@ -133,6 +134,18 @@ class SwitchTemplate(SwitchEntity):
     async def async_added_to_hass(self):
         """Register callbacks."""
 
+        if self._template is None:
+
+            # restore state after startup
+            await super().async_added_to_hass()
+            state = await self.async_get_last_state()
+            if state:
+                self._state = state.state == STATE_ON
+
+            # no need to listen for events
+            return
+
+        # set up event listening
         @callback
         def template_switch_state_listener(entity, old_state, new_state):
             """Handle target device state changes."""
@@ -184,13 +197,21 @@ class SwitchTemplate(SwitchEntity):
     async def async_turn_on(self, **kwargs):
         """Fire the on action."""
         await self._on_script.async_run(context=self._context)
+        if self._template is None:
+            self._state = True
+            self.async_write_ha_state()
 
     async def async_turn_off(self, **kwargs):
         """Fire the off action."""
         await self._off_script.async_run(context=self._context)
+        if self._template is None:
+            self._state = False
+            self.async_write_ha_state()
 
     async def async_update(self):
         """Update the state from the template."""
+        if self._template is None:
+            return
         try:
             state = self._template.async_render().lower()
 
@@ -243,3 +264,8 @@ class SwitchTemplate(SwitchEntity):
                         self._name,
                         ex,
                     )
+
+    @property
+    def assumed_state(self):
+        """State is assumed, if no template given."""
+        return self._template is None

tests/components/template/test_switch.py

@@ -1,40 +1,32 @@
 """The tests for the  Template switch platform."""
+
+import pytest
+
 from homeassistant import setup
 from homeassistant.const import STATE_OFF, STATE_ON, STATE_UNAVAILABLE
-from homeassistant.core import callback
+from homeassistant.core import CoreState, State
+from homeassistant.setup import async_setup_component
 
-from tests.common import assert_setup_component, get_test_home_assistant
+from tests.common import (
+    assert_setup_component,
+    async_mock_service,
+    mock_component,
+    mock_restore_cache,
+)
 from tests.components.switch import common
 
 
-class TestTemplateSwitch:
+@pytest.fixture
-    """Test the Template switch."""
+def calls(hass):
+    """Track calls to a mock service."""
+    return async_mock_service(hass, "test", "automation")
 
-    hass = None
-    calls = None
-    # pylint: disable=invalid-name
 
-    def setup_method(self, method):
+async def test_template_state_text(hass):
-        """Set up things to be run when tests are started."""
-        self.hass = get_test_home_assistant()
-        self.calls = []
-
-        @callback
-        def record_call(service):
-            """Track function calls.."""
-            self.calls.append(service)
-
-        self.hass.services.register("test", "automation", record_call)
-
-    def teardown_method(self, method):
-        """Stop everything that was started."""
-        self.hass.stop()
-
-    def test_template_state_text(self):
     """Test the state text of a template."""
     with assert_setup_component(1, "switch"):
-            assert setup.setup_component(
+        assert await async_setup_component(
-                self.hass,
+            hass,
             "switch",
             {
                 "switch": {
@@ -56,27 +48,28 @@ class TestTemplateSwitch:
             },
         )
 
-        self.hass.block_till_done()
+    await hass.async_block_till_done()
-        self.hass.start()
+    await hass.async_start()
-        self.hass.block_till_done()
+    await hass.async_block_till_done()
 
-        state = self.hass.states.set("switch.test_state", STATE_ON)
+    hass.states.async_set("switch.test_state", STATE_ON)
-        self.hass.block_till_done()
+    await hass.async_block_till_done()
 
-        state = self.hass.states.get("switch.test_template_switch")
+    state = hass.states.get("switch.test_template_switch")
     assert state.state == STATE_ON
 
-        state = self.hass.states.set("switch.test_state", STATE_OFF)
+    hass.states.async_set("switch.test_state", STATE_OFF)
-        self.hass.block_till_done()
+    await hass.async_block_till_done()
 
-        state = self.hass.states.get("switch.test_template_switch")
+    state = hass.states.get("switch.test_template_switch")
     assert state.state == STATE_OFF
 
-    def test_template_state_boolean_on(self):
+
+async def test_template_state_boolean_on(hass):
     """Test the setting of the state with boolean on."""
     with assert_setup_component(1, "switch"):
-            assert setup.setup_component(
+        assert await async_setup_component(
-                self.hass,
+            hass,
             "switch",
             {
                 "switch": {
@@ -98,18 +91,19 @@ class TestTemplateSwitch:
             },
         )
 
-        self.hass.block_till_done()
+    await hass.async_block_till_done()
-        self.hass.start()
+    await hass.async_start()
-        self.hass.block_till_done()
+    await hass.async_block_till_done()
 
-        state = self.hass.states.get("switch.test_template_switch")
+    state = hass.states.get("switch.test_template_switch")
     assert state.state == STATE_ON
 
-    def test_template_state_boolean_off(self):
+
+async def test_template_state_boolean_off(hass):
     """Test the setting of the state with off."""
     with assert_setup_component(1, "switch"):
-            assert setup.setup_component(
+        assert await async_setup_component(
-                self.hass,
+            hass,
             "switch",
             {
                 "switch": {
@@ -131,18 +125,19 @@ class TestTemplateSwitch:
             },
         )
 
-        self.hass.block_till_done()
+    await hass.async_block_till_done()
-        self.hass.start()
+    await hass.async_start()
-        self.hass.block_till_done()
+    await hass.async_block_till_done()
 
-        state = self.hass.states.get("switch.test_template_switch")
+    state = hass.states.get("switch.test_template_switch")
     assert state.state == STATE_OFF
 
-    def test_icon_template(self):
+
+async def test_icon_template(hass):
     """Test icon template."""
     with assert_setup_component(1, "switch"):
-            assert setup.setup_component(
+        assert await async_setup_component(
-                self.hass,
+            hass,
             "switch",
             {
                 "switch": {
@@ -167,24 +162,25 @@ class TestTemplateSwitch:
             },
         )
 
-        self.hass.block_till_done()
+    await hass.async_block_till_done()
-        self.hass.start()
+    await hass.async_start()
-        self.hass.block_till_done()
+    await hass.async_block_till_done()
 
-        state = self.hass.states.get("switch.test_template_switch")
+    state = hass.states.get("switch.test_template_switch")
     assert state.attributes.get("icon") == ""
 
-        state = self.hass.states.set("switch.test_state", STATE_ON)
+    hass.states.async_set("switch.test_state", STATE_ON)
-        self.hass.block_till_done()
+    await hass.async_block_till_done()
 
-        state = self.hass.states.get("switch.test_template_switch")
+    state = hass.states.get("switch.test_template_switch")
     assert state.attributes["icon"] == "mdi:check"
 
-    def test_entity_picture_template(self):
+
+async def test_entity_picture_template(hass):
     """Test entity_picture template."""
     with assert_setup_component(1, "switch"):
-            assert setup.setup_component(
+        assert await async_setup_component(
-                self.hass,
+            hass,
             "switch",
             {
                 "switch": {
@@ -209,24 +205,25 @@ class TestTemplateSwitch:
             },
         )
 
-        self.hass.block_till_done()
+    await hass.async_block_till_done()
-        self.hass.start()
+    await hass.async_start()
-        self.hass.block_till_done()
+    await hass.async_block_till_done()
 
-        state = self.hass.states.get("switch.test_template_switch")
+    state = hass.states.get("switch.test_template_switch")
     assert state.attributes.get("entity_picture") == ""
 
-        state = self.hass.states.set("switch.test_state", STATE_ON)
+    hass.states.async_set("switch.test_state", STATE_ON)
-        self.hass.block_till_done()
+    await hass.async_block_till_done()
 
-        state = self.hass.states.get("switch.test_template_switch")
+    state = hass.states.get("switch.test_template_switch")
     assert state.attributes["entity_picture"] == "/local/switch.png"
 
-    def test_template_syntax_error(self):
+
+async def test_template_syntax_error(hass):
     """Test templating syntax error."""
     with assert_setup_component(0, "switch"):
-            assert setup.setup_component(
+        assert await async_setup_component(
-                self.hass,
+            hass,
             "switch",
             {
                 "switch": {
@@ -248,17 +245,18 @@ class TestTemplateSwitch:
             },
         )
 
-        self.hass.block_till_done()
+    await hass.async_block_till_done()
-        self.hass.start()
+    await hass.async_start()
-        self.hass.block_till_done()
+    await hass.async_block_till_done()
 
-        assert self.hass.states.all() == []
+    assert hass.states.async_all() == []
 
-    def test_invalid_name_does_not_create(self):
+
+async def test_invalid_name_does_not_create(hass):
     """Test invalid name."""
     with assert_setup_component(0, "switch"):
-            assert setup.setup_component(
+        assert await async_setup_component(
-                self.hass,
+            hass,
             "switch",
             {
                 "switch": {
@@ -280,17 +278,18 @@ class TestTemplateSwitch:
             },
         )
 
-        self.hass.block_till_done()
+    await hass.async_block_till_done()
-        self.hass.start()
+    await hass.async_start()
-        self.hass.block_till_done()
+    await hass.async_block_till_done()
 
-        assert self.hass.states.all() == []
+    assert hass.states.async_all() == []
 
-    def test_invalid_switch_does_not_create(self):
+
+async def test_invalid_switch_does_not_create(hass):
     """Test invalid switch."""
     with assert_setup_component(0, "switch"):
-            assert setup.setup_component(
+        assert await async_setup_component(
-                self.hass,
+            hass,
             "switch",
             {
                 "switch": {
@@ -300,62 +299,32 @@ class TestTemplateSwitch:
             },
         )
 
-        self.hass.block_till_done()
+    await hass.async_block_till_done()
-        self.hass.start()
+    await hass.async_start()
-        self.hass.block_till_done()
+    await hass.async_block_till_done()
 
-        assert self.hass.states.all() == []
+    assert hass.states.async_all() == []
 
-    def test_no_switches_does_not_create(self):
+
+async def test_no_switches_does_not_create(hass):
     """Test if there are no switches no creation."""
     with assert_setup_component(0, "switch"):
-            assert setup.setup_component(
+        assert await async_setup_component(
-                self.hass, "switch", {"switch": {"platform": "template"}}
+            hass, "switch", {"switch": {"platform": "template"}}
         )
 
-        self.hass.block_till_done()
+    await hass.async_block_till_done()
-        self.hass.start()
+    await hass.async_start()
-        self.hass.block_till_done()
+    await hass.async_block_till_done()
 
-        assert self.hass.states.all() == []
+    assert hass.states.async_all() == []
 
-    def test_missing_template_does_not_create(self):
-        """Test missing template."""
-        with assert_setup_component(0, "switch"):
-            assert setup.setup_component(
-                self.hass,
-                "switch",
-                {
-                    "switch": {
-                        "platform": "template",
-                        "switches": {
-                            "test_template_switch": {
-                                "not_value_template": "{{ states.switch.test_state.state }}",
-                                "turn_on": {
-                                    "service": "switch.turn_on",
-                                    "entity_id": "switch.test_state",
-                                },
-                                "turn_off": {
-                                    "service": "switch.turn_off",
-                                    "entity_id": "switch.test_state",
-                                },
-                            }
-                        },
-                    }
-                },
-            )
 
-        self.hass.block_till_done()
+async def test_missing_on_does_not_create(hass):
-        self.hass.start()
-        self.hass.block_till_done()
-
-        assert self.hass.states.all() == []
-
-    def test_missing_on_does_not_create(self):
     """Test missing on."""
     with assert_setup_component(0, "switch"):
-            assert setup.setup_component(
+        assert await async_setup_component(
-                self.hass,
+            hass,
             "switch",
             {
                 "switch": {
@@ -377,17 +346,18 @@ class TestTemplateSwitch:
             },
         )
 
-        self.hass.block_till_done()
+    await hass.async_block_till_done()
-        self.hass.start()
+    await hass.async_start()
-        self.hass.block_till_done()
+    await hass.async_block_till_done()
 
-        assert self.hass.states.all() == []
+    assert hass.states.async_all() == []
 
-    def test_missing_off_does_not_create(self):
+
+async def test_missing_off_does_not_create(hass):
     """Test missing off."""
     with assert_setup_component(0, "switch"):
-            assert setup.setup_component(
+        assert await async_setup_component(
-                self.hass,
+            hass,
             "switch",
             {
                 "switch": {
@@ -409,16 +379,17 @@ class TestTemplateSwitch:
             },
         )
 
-        self.hass.block_till_done()
+    await hass.async_block_till_done()
-        self.hass.start()
+    await hass.async_start()
-        self.hass.block_till_done()
+    await hass.async_block_till_done()
 
-        assert self.hass.states.all() == []
+    assert hass.states.async_all() == []
 
-    def test_on_action(self):
+
+async def test_on_action(hass, calls):
     """Test on action."""
-        assert setup.setup_component(
+    assert await async_setup_component(
-            self.hass,
+        hass,
         "switch",
         {
             "switch": {
@@ -437,25 +408,64 @@ class TestTemplateSwitch:
         },
     )
 
-        self.hass.block_till_done()
+    await hass.async_block_till_done()
-        self.hass.start()
+    await hass.async_start()
-        self.hass.block_till_done()
+    await hass.async_block_till_done()
 
-        self.hass.states.set("switch.test_state", STATE_OFF)
+    hass.states.async_set("switch.test_state", STATE_OFF)
-        self.hass.block_till_done()
+    await hass.async_block_till_done()
 
-        state = self.hass.states.get("switch.test_template_switch")
+    state = hass.states.get("switch.test_template_switch")
     assert state.state == STATE_OFF
 
-        common.turn_on(self.hass, "switch.test_template_switch")
+    await common.async_turn_on(hass, "switch.test_template_switch")
-        self.hass.block_till_done()
+    await hass.async_block_till_done()
 
-        assert len(self.calls) == 1
+    assert len(calls) == 1
 
-    def test_off_action(self):
+
+async def test_on_action_optimistic(hass, calls):
+    """Test on action in optimistic mode."""
+    assert await async_setup_component(
+        hass,
+        "switch",
+        {
+            "switch": {
+                "platform": "template",
+                "switches": {
+                    "test_template_switch": {
+                        "turn_on": {"service": "test.automation"},
+                        "turn_off": {
+                            "service": "switch.turn_off",
+                            "entity_id": "switch.test_state",
+                        },
+                    }
+                },
+            }
+        },
+    )
+
+    await hass.async_start()
+    await hass.async_block_till_done()
+
+    hass.states.async_set("switch.test_template_switch", STATE_OFF)
+    await hass.async_block_till_done()
+
+    state = hass.states.get("switch.test_template_switch")
+    assert state.state == STATE_OFF
+
+    await common.async_turn_on(hass, "switch.test_template_switch")
+    await hass.async_block_till_done()
+
+    state = hass.states.get("switch.test_template_switch")
+    assert len(calls) == 1
+    assert state.state == STATE_ON
+
+
+async def test_off_action(hass, calls):
     """Test off action."""
-        assert setup.setup_component(
+    assert await async_setup_component(
-            self.hass,
+        hass,
         "switch",
         {
             "switch": {
@@ -474,20 +484,97 @@ class TestTemplateSwitch:
         },
     )
 
-        self.hass.block_till_done()
+    await hass.async_block_till_done()
-        self.hass.start()
+    await hass.async_start()
-        self.hass.block_till_done()
+    await hass.async_block_till_done()
 
-        self.hass.states.set("switch.test_state", STATE_ON)
+    hass.states.async_set("switch.test_state", STATE_ON)
-        self.hass.block_till_done()
+    await hass.async_block_till_done()
 
-        state = self.hass.states.get("switch.test_template_switch")
+    state = hass.states.get("switch.test_template_switch")
     assert state.state == STATE_ON
 
-        common.turn_off(self.hass, "switch.test_template_switch")
+    await common.async_turn_off(hass, "switch.test_template_switch")
-        self.hass.block_till_done()
+    await hass.async_block_till_done()
 
-        assert len(self.calls) == 1
+    assert len(calls) == 1
+
+
+async def test_off_action_optimistic(hass, calls):
+    """Test off action in optimistic mode."""
+    assert await async_setup_component(
+        hass,
+        "switch",
+        {
+            "switch": {
+                "platform": "template",
+                "switches": {
+                    "test_template_switch": {
+                        "turn_off": {"service": "test.automation"},
+                        "turn_on": {
+                            "service": "switch.turn_on",
+                            "entity_id": "switch.test_state",
+                        },
+                    }
+                },
+            }
+        },
+    )
+
+    await hass.async_start()
+    await hass.async_block_till_done()
+
+    hass.states.async_set("switch.test_template_switch", STATE_ON)
+    await hass.async_block_till_done()
+
+    state = hass.states.get("switch.test_template_switch")
+    assert state.state == STATE_ON
+
+    await common.async_turn_off(hass, "switch.test_template_switch")
+    await hass.async_block_till_done()
+
+    state = hass.states.get("switch.test_template_switch")
+    assert len(calls) == 1
+    assert state.state == STATE_OFF
+
+
+async def test_restore_state(hass):
+    """Test state restoration."""
+    mock_restore_cache(
+        hass, (State("switch.s1", STATE_ON), State("switch.s2", STATE_OFF),),
+    )
+
+    hass.state = CoreState.starting
+    mock_component(hass, "recorder")
+
+    await async_setup_component(
+        hass,
+        "switch",
+        {
+            "switch": {
+                "platform": "template",
+                "switches": {
+                    "s1": {
+                        "turn_on": {"service": "test.automation"},
+                        "turn_off": {"service": "test.automation"},
+                    },
+                    "s2": {
+                        "turn_on": {"service": "test.automation"},
+                        "turn_off": {"service": "test.automation"},
+                    },
+                },
+            }
+        },
+    )
+    await hass.async_block_till_done()
+
+    state = hass.states.get("switch.s1")
+    assert state
+    assert state.state == STATE_ON
+
+    state = hass.states.get("switch.s2")
+    assert state
+    assert state.state == STATE_OFF
 
 
 async def test_available_template_with_entities(hass):