Add optimistic option to switch yaml (#149402)

homeassistant/components/template/switch.py

@@ -38,6 +38,7 @@ from homeassistant.helpers.typing import ConfigType, DiscoveryInfoType
 
 from . import TriggerUpdateCoordinator
 from .const import CONF_TURN_OFF, CONF_TURN_ON, DOMAIN
+from .entity import AbstractTemplateEntity
 from .helpers import (
     async_setup_template_entry,
     async_setup_template_platform,
@@ -46,6 +47,7 @@ from .helpers import (
 from .template_entity import (
     TEMPLATE_ENTITY_COMMON_CONFIG_ENTRY_SCHEMA,
     TEMPLATE_ENTITY_COMMON_SCHEMA_LEGACY,
+    TEMPLATE_ENTITY_OPTIMISTIC_SCHEMA,
     TemplateEntity,
     make_template_entity_common_modern_schema,
 )
@@ -68,8 +70,8 @@ SWITCH_COMMON_SCHEMA = vol.Schema(
 )
 
 SWITCH_YAML_SCHEMA = SWITCH_COMMON_SCHEMA.extend(
-    make_template_entity_common_modern_schema(DEFAULT_NAME).schema
+    TEMPLATE_ENTITY_OPTIMISTIC_SCHEMA
-)
+).extend(make_template_entity_common_modern_schema(DEFAULT_NAME).schema)
 
 SWITCH_LEGACY_YAML_SCHEMA = vol.All(
     cv.deprecated(ATTR_ENTITY_ID),
@@ -145,11 +147,38 @@ def async_create_preview_switch(
     )
 
 
-class StateSwitchEntity(TemplateEntity, SwitchEntity, RestoreEntity):
+class AbstractTemplateSwitch(AbstractTemplateEntity, SwitchEntity, RestoreEntity):
+    """Representation of a template switch features."""
+
+    _entity_id_format = ENTITY_ID_FORMAT
+    _optimistic_entity = True
+
+    # The super init is not called because TemplateEntity and TriggerEntity will call AbstractTemplateEntity.__init__.
+    # This ensures that the __init__ on AbstractTemplateEntity is not called twice.
+    def __init__(self, config: dict[str, Any]) -> None:  # pylint: disable=super-init-not-called
+        """Initialize the features."""
+
+    async def async_turn_on(self, **kwargs: Any) -> None:
+        """Fire the on action."""
+        if on_script := self._action_scripts.get(CONF_TURN_ON):
+            await self.async_run_script(on_script, context=self._context)
+        if self._attr_assumed_state:
+            self._attr_is_on = True
+            self.async_write_ha_state()
+
+    async def async_turn_off(self, **kwargs: Any) -> None:
+        """Fire the off action."""
+        if off_script := self._action_scripts.get(CONF_TURN_OFF):
+            await self.async_run_script(off_script, context=self._context)
+        if self._attr_assumed_state:
+            self._attr_is_on = False
+            self.async_write_ha_state()
+
+
+class StateSwitchEntity(TemplateEntity, AbstractTemplateSwitch):
     """Representation of a Template switch."""
 
     _attr_should_poll = False
-    _entity_id_format = ENTITY_ID_FORMAT
 
     def __init__(
         self,
@@ -158,12 +187,12 @@ class StateSwitchEntity(TemplateEntity, SwitchEntity, RestoreEntity):
         unique_id: str | None,
     ) -> None:
         """Initialize the Template switch."""
-        super().__init__(hass, config, unique_id)
+        TemplateEntity.__init__(self, hass, config, unique_id)
+        AbstractTemplateSwitch.__init__(self, config)
 
         name = self._attr_name
         if TYPE_CHECKING:
             assert name is not None
-        self._template = config.get(CONF_STATE)
 
         # Scripts can be an empty list, therefore we need to check for None
         if (on_action := config.get(CONF_TURN_ON)) is not None:
@@ -171,25 +200,22 @@ class StateSwitchEntity(TemplateEntity, SwitchEntity, RestoreEntity):
         if (off_action := config.get(CONF_TURN_OFF)) is not None:
             self.add_script(CONF_TURN_OFF, off_action, name, DOMAIN)
 
-        self._state: bool | None = False
-        self._attr_assumed_state = self._template is None
-
     @callback
     def _update_state(self, result):
         super()._update_state(result)
         if isinstance(result, TemplateError):
-            self._state = None
+            self._attr_is_on = None
             return
 
         if isinstance(result, bool):
-            self._state = result
+            self._attr_is_on = result
             return
 
         if isinstance(result, str):
-            self._state = result.lower() in ("true", STATE_ON)
+            self._attr_is_on = result.lower() in ("true", STATE_ON)
             return
 
-        self._state = False
+        self._attr_is_on = False
 
     async def async_added_to_hass(self) -> None:
         """Register callbacks."""
@@ -197,7 +223,7 @@ class StateSwitchEntity(TemplateEntity, SwitchEntity, RestoreEntity):
             # restore state after startup
             await super().async_added_to_hass()
             if state := await self.async_get_last_state():
-                self._state = state.state == STATE_ON
+                self._attr_is_on = state.state == STATE_ON
         await super().async_added_to_hass()
 
     @callback
@@ -205,37 +231,15 @@ class StateSwitchEntity(TemplateEntity, SwitchEntity, RestoreEntity):
         """Set up templates."""
         if self._template is not None:
             self.add_template_attribute(
-                "_state", self._template, None, self._update_state
+                "_attr_is_on", self._template, None, self._update_state
             )
 
         super()._async_setup_templates()
 
-    @property
-    def is_on(self) -> bool | None:
-        """Return true if device is on."""
-        return self._state
 
-    async def async_turn_on(self, **kwargs: Any) -> None:
+class TriggerSwitchEntity(TriggerEntity, AbstractTemplateSwitch):
-        """Fire the on action."""
-        if on_script := self._action_scripts.get(CONF_TURN_ON):
-            await self.async_run_script(on_script, context=self._context)
-        if self._template is None:
-            self._state = True
-            self.async_write_ha_state()
-
-    async def async_turn_off(self, **kwargs: Any) -> None:
-        """Fire the off action."""
-        if off_script := self._action_scripts.get(CONF_TURN_OFF):
-            await self.async_run_script(off_script, context=self._context)
-        if self._template is None:
-            self._state = False
-            self.async_write_ha_state()
-
-
-class TriggerSwitchEntity(TriggerEntity, SwitchEntity, RestoreEntity):
     """Switch entity based on trigger data."""
 
-    _entity_id_format = ENTITY_ID_FORMAT
     domain = SWITCH_DOMAIN
 
     def __init__(
@@ -245,17 +249,16 @@ class TriggerSwitchEntity(TriggerEntity, SwitchEntity, RestoreEntity):
         config: ConfigType,
     ) -> None:
         """Initialize the entity."""
-        super().__init__(hass, coordinator, config)
+        TriggerEntity.__init__(self, hass, coordinator, config)
+        AbstractTemplateSwitch.__init__(self, config)
 
         name = self._rendered.get(CONF_NAME, DEFAULT_NAME)
-        self._template = config.get(CONF_STATE)
         if on_action := config.get(CONF_TURN_ON):
             self.add_script(CONF_TURN_ON, on_action, name, DOMAIN)
         if off_action := config.get(CONF_TURN_OFF):
             self.add_script(CONF_TURN_OFF, off_action, name, DOMAIN)
 
-        self._attr_assumed_state = self._template is None
+        if CONF_STATE in config:
-        if not self._attr_assumed_state:
             self._to_render_simple.append(CONF_STATE)
             self._parse_result.add(CONF_STATE)
 
@@ -281,28 +284,15 @@ class TriggerSwitchEntity(TriggerEntity, SwitchEntity, RestoreEntity):
             self.async_write_ha_state()
             return
 
-        if not self._attr_assumed_state:
+        write_ha_state = False
-            raw = self._rendered.get(CONF_STATE)
+        if (state := self._rendered.get(CONF_STATE)) is not None:
-            self._attr_is_on = template.result_as_boolean(raw)
+            self._attr_is_on = template.result_as_boolean(state)
+            write_ha_state = True
 
-            self.async_write_ha_state()
+        elif len(self._rendered) > 0:
-        elif self._attr_assumed_state and len(self._rendered) > 0:
             # In case name, icon, or friendly name have a template but
             # states does not
-            self.async_write_ha_state()
+            write_ha_state = True
 
-    async def async_turn_on(self, **kwargs: Any) -> None:
+        if write_ha_state:
-        """Fire the on action."""
-        if on_script := self._action_scripts.get(CONF_TURN_ON):
-            await self.async_run_script(on_script, context=self._context)
-        if self._template is None:
-            self._attr_is_on = True
-            self.async_write_ha_state()
-
-    async def async_turn_off(self, **kwargs: Any) -> None:
-        """Fire the off action."""
-        if off_script := self._action_scripts.get(CONF_TURN_OFF):
-            await self.async_run_script(off_script, context=self._context)
-        if self._template is None:
-            self._attr_is_on = False
             self.async_write_ha_state()

tests/components/template/test_switch.py

@@ -34,8 +34,13 @@ TEST_ENTITY_ID = f"switch.{TEST_OBJECT_ID}"
 TEST_STATE_ENTITY_ID = "switch.test_state"
 
 TEST_EVENT_TRIGGER = {
-    "trigger": {"platform": "event", "event_type": "test_event"},
+    "triggers": [
-    "variables": {"type": "{{ trigger.event.data.type }}"},
+        {"trigger": "event", "event_type": "test_event"},
+        {"trigger": "state", "entity_id": [TEST_STATE_ENTITY_ID]},
+    ],
+    "variables": {
+        "type": "{{ trigger.event.data.type if trigger.event is defined else trigger.entity_id }}"
+    },
     "action": [{"event": "action_event", "event_data": {"type": "{{ type }}"}}],
 }
 
@@ -1211,3 +1216,54 @@ async def test_empty_action_config(hass: HomeAssistant, setup_switch) -> None:
 
     state = hass.states.get(TEST_ENTITY_ID)
     assert state.state == STATE_OFF
+
+
+@pytest.mark.parametrize(
+    ("count", "switch_config"),
+    [
+        (
+            1,
+            {
+                "name": TEST_OBJECT_ID,
+                "state": "{{ is_state('switch.test_state', 'on') }}",
+                "turn_on": [],
+                "turn_off": [],
+                "optimistic": True,
+            },
+        )
+    ],
+)
+@pytest.mark.parametrize(
+    "style",
+    [
+        ConfigurationStyle.MODERN,
+        ConfigurationStyle.TRIGGER,
+    ],
+)
+@pytest.mark.usefixtures("setup_switch")
+async def test_optimistic_option(hass: HomeAssistant) -> None:
+    """Test optimistic yaml option."""
+    hass.states.async_set(TEST_STATE_ENTITY_ID, STATE_OFF)
+    await hass.async_block_till_done()
+
+    state = hass.states.get(TEST_ENTITY_ID)
+    assert state.state == STATE_OFF
+
+    await hass.services.async_call(
+        switch.DOMAIN,
+        "turn_on",
+        {"entity_id": TEST_ENTITY_ID},
+        blocking=True,
+    )
+
+    state = hass.states.get(TEST_ENTITY_ID)
+    assert state.state == STATE_ON
+
+    hass.states.async_set(TEST_STATE_ENTITY_ID, STATE_ON)
+    await hass.async_block_till_done()
+
+    hass.states.async_set(TEST_STATE_ENTITY_ID, STATE_OFF)
+    await hass.async_block_till_done()
+
+    state = hass.states.get(TEST_ENTITY_ID)
+    assert state.state == STATE_OFF