Add assumed optimistic state to template select (home-assistant#148513)

Author: Petro31

homeassistant/components/template/select.py

@@ -3,7 +3,7 @@
 from __future__ import annotations
 
 import logging
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 import voluptuous as vol
 
@@ -32,6 +32,7 @@
 
 from . import TriggerUpdateCoordinator
 from .const import DOMAIN
+from .entity import AbstractTemplateEntity
 from .template_entity import TemplateEntity, make_template_entity_common_modern_schema
 from .trigger_entity import TriggerEntity
 
@@ -45,7 +46,7 @@
 
 SELECT_SCHEMA = vol.Schema(
     {
-        vol.Required(CONF_STATE): cv.template,
+        vol.Optional(CONF_STATE): cv.template,
         vol.Required(CONF_SELECT_OPTION): cv.SCRIPT_SCHEMA,
         vol.Required(ATTR_OPTIONS): cv.template,
         vol.Optional(CONF_OPTIMISTIC, default=DEFAULT_OPTIMISTIC): cv.boolean,
@@ -116,7 +117,37 @@ async def async_setup_entry(
     async_add_entities([TemplateSelect(hass, validated_config, config_entry.entry_id)])
 
 
-class TemplateSelect(TemplateEntity, SelectEntity):
+class AbstractTemplateSelect(AbstractTemplateEntity, SelectEntity):
+    """Representation of a template select features."""
+
+    # 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."""
+        self._template = config.get(CONF_STATE)
+
+        self._options_template = config[ATTR_OPTIONS]
+
+        self._attr_assumed_state = self._optimistic = (
+            self._template is None or config.get(CONF_OPTIMISTIC, DEFAULT_OPTIMISTIC)
+        )
+        self._attr_options = []
+        self._attr_current_option = None
+
+    async def async_select_option(self, option: str) -> None:
+        """Change the selected option."""
+        if self._optimistic:
+            self._attr_current_option = option
+            self.async_write_ha_state()
+        if select_option := self._action_scripts.get(CONF_SELECT_OPTION):
+            await self.async_run_script(
+                select_option,
+                run_variables={ATTR_OPTION: option},
+                context=self._context,
+            )
+
+
+class TemplateSelect(TemplateEntity, AbstractTemplateSelect):
     """Representation of a template select."""
 
     _attr_should_poll = False
@@ -128,16 +159,16 @@ def __init__(
         unique_id: str | None,
     ) -> None:
         """Initialize the select."""
-        super().__init__(hass, config=config, unique_id=unique_id)
-        assert self._attr_name is not None
-        self._value_template = config[CONF_STATE]
-        # Scripts can be an empty list, therefore we need to check for None
+        TemplateEntity.__init__(self, hass, config=config, unique_id=unique_id)
+        AbstractTemplateSelect.__init__(self, config)
+
+        name = self._attr_name
+        if TYPE_CHECKING:
+            assert name is not None
+
         if (select_option := config.get(CONF_SELECT_OPTION)) is not None:
-            self.add_script(CONF_SELECT_OPTION, select_option, self._attr_name, DOMAIN)
-        self._options_template = config[ATTR_OPTIONS]
-        self._attr_assumed_state = self._optimistic = config.get(CONF_OPTIMISTIC, False)
-        self._attr_options = []
-        self._attr_current_option = None
+            self.add_script(CONF_SELECT_OPTION, select_option, name, DOMAIN)
+
         self._attr_device_info = async_device_info_to_link_from_device_id(
             hass,
             config.get(CONF_DEVICE_ID),
@@ -146,12 +177,13 @@ def __init__(
     @callback
     def _async_setup_templates(self) -> None:
         """Set up templates."""
-        self.add_template_attribute(
-            "_attr_current_option",
-            self._value_template,
-            validator=cv.string,
-            none_on_template_error=True,
-        )
+        if self._template is not None:
+            self.add_template_attribute(
+                "_attr_current_option",
+                self._template,
+                validator=cv.string,
+                none_on_template_error=True,
+            )
         self.add_template_attribute(
             "_attr_options",
             self._options_template,
@@ -160,24 +192,11 @@ def _async_setup_templates(self) -> None:
         )
         super()._async_setup_templates()
 
-    async def async_select_option(self, option: str) -> None:
-        """Change the selected option."""
-        if self._optimistic:
-            self._attr_current_option = option
-            self.async_write_ha_state()
-        if select_option := self._action_scripts.get(CONF_SELECT_OPTION):
-            await self.async_run_script(
-                select_option,
-                run_variables={ATTR_OPTION: option},
-                context=self._context,
-            )
-
 
-class TriggerSelectEntity(TriggerEntity, SelectEntity):
+class TriggerSelectEntity(TriggerEntity, AbstractTemplateSelect):
     """Select entity based on trigger data."""
 
     domain = SELECT_DOMAIN
-    extra_template_keys = (CONF_STATE,)
     extra_template_keys_complex = (ATTR_OPTIONS,)
 
     def __init__(
@@ -187,7 +206,12 @@ def __init__(
         config: dict,
     ) -> None:
         """Initialize the entity."""
-        super().__init__(hass, coordinator, config)
+        TriggerEntity.__init__(self, hass, coordinator, config)
+        AbstractTemplateSelect.__init__(self, config)
+
+        if CONF_STATE in config:
+            self._to_render_simple.append(CONF_STATE)
+
         # Scripts can be an empty list, therefore we need to check for None
         if (select_option := config.get(CONF_SELECT_OPTION)) is not None:
             self.add_script(
@@ -197,24 +221,26 @@ def __init__(
                 DOMAIN,
             )
 
-    @property
-    def current_option(self) -> str | None:
-        """Return the currently selected option."""
-        return self._rendered.get(CONF_STATE)
+    def _handle_coordinator_update(self):
+        """Handle updated data from the coordinator."""
+        self._process_data()
 
-    @property
-    def options(self) -> list[str]:
-        """Return the list of available options."""
-        return self._rendered.get(ATTR_OPTIONS, [])
+        if not self.available:
+            self.async_write_ha_state()
+            return
 
-    async def async_select_option(self, option: str) -> None:
-        """Change the selected option."""
-        if self._config[CONF_OPTIMISTIC]:
-            self._attr_current_option = option
+        write_ha_state = False
+        if (options := self._rendered.get(ATTR_OPTIONS)) is not None:
+            self._attr_options = vol.All(cv.ensure_list, [cv.string])(options)
+            write_ha_state = True
+
+        if (state := self._rendered.get(CONF_STATE)) is not None:
+            self._attr_current_option = cv.string(state)
+            write_ha_state = True
+
+        if len(self._rendered) > 0:
+            # In case any non optimistic template
+            write_ha_state = True
+
+        if write_ha_state:
             self.async_write_ha_state()
-        if select_option := self._action_scripts.get(CONF_SELECT_OPTION):
-            await self.async_run_script(
-                select_option,
-                run_variables={ATTR_OPTION: option},
-                context=self._context,
-            )

tests/components/template/test_select.py

@@ -28,6 +28,7 @@
     ATTR_ICON,
     CONF_ENTITY_ID,
     CONF_ICON,
+    STATE_UNAVAILABLE,
     STATE_UNKNOWN,
 )
 from homeassistant.core import Context, HomeAssistant, ServiceCall
@@ -43,11 +44,15 @@
 # Represent for select's current_option
 _OPTION_INPUT_SELECT = "input_select.option"
 TEST_STATE_ENTITY_ID = "select.test_state"
-
+TEST_AVAILABILITY_ENTITY_ID = "binary_sensor.test_availability"
 TEST_STATE_TRIGGER = {
     "trigger": {
         "trigger": "state",
-        "entity_id": [_OPTION_INPUT_SELECT, TEST_STATE_ENTITY_ID],
+        "entity_id": [
+            _OPTION_INPUT_SELECT,
+            TEST_STATE_ENTITY_ID,
+            TEST_AVAILABILITY_ENTITY_ID,
+        ],
     },
     "variables": {"triggering_entity": "{{ trigger.entity_id }}"},
     "action": [
@@ -201,20 +206,6 @@ async def test_multiple_configs(hass: HomeAssistant) -> None:
 
 async def test_missing_required_keys(hass: HomeAssistant) -> None:
     """Test: missing required fields will fail."""
-    with assert_setup_component(0, "template"):
-        assert await setup.async_setup_component(
-            hass,
-            "template",
-            {
-                "template": {
-                    "select": {
-                        "select_option": {"service": "script.select_option"},
-                        "options": "{{ ['a', 'b'] }}",
-                    }
-                }
-            },
-        )
-
     with assert_setup_component(0, "select"):
         assert await setup.async_setup_component(
             hass,
@@ -559,3 +550,98 @@ async def test_empty_action_config(hass: HomeAssistant, setup_select) -> None:
 
     state = hass.states.get(_TEST_SELECT)
     assert state.state == "a"
+
+
+@pytest.mark.parametrize(
+    ("count", "select_config"),
+    [
+        (
+            1,
+            {
+                "options": "{{ ['test', 'yes', 'no'] }}",
+                "select_option": [],
+            },
+        )
+    ],
+)
+@pytest.mark.parametrize(
+    "style",
+    [ConfigurationStyle.MODERN, ConfigurationStyle.TRIGGER],
+)
+@pytest.mark.usefixtures("setup_select")
+async def test_optimistic(hass: HomeAssistant) -> None:
+    """Test configuration with optimistic state."""
+
+    state = hass.states.get(_TEST_SELECT)
+    assert state.state == STATE_UNKNOWN
+
+    # Ensure Trigger template entities update.
+    hass.states.async_set(TEST_STATE_ENTITY_ID, "anything")
+    await hass.async_block_till_done()
+
+    await hass.services.async_call(
+        select.DOMAIN,
+        select.SERVICE_SELECT_OPTION,
+        {ATTR_ENTITY_ID: _TEST_SELECT, "option": "test"},
+        blocking=True,
+    )
+
+    state = hass.states.get(_TEST_SELECT)
+    assert state.state == "test"
+
+    await hass.services.async_call(
+        select.DOMAIN,
+        select.SERVICE_SELECT_OPTION,
+        {ATTR_ENTITY_ID: _TEST_SELECT, "option": "yes"},
+        blocking=True,
+    )
+
+    state = hass.states.get(_TEST_SELECT)
+    assert state.state == "yes"
+
+
+@pytest.mark.parametrize(
+    ("count", "select_config"),
+    [
+        (
+            1,
+            {
+                "options": "{{ ['test', 'yes', 'no'] }}",
+                "select_option": [],
+                "state": "{{ states('select.test_state') }}",
+                "availability": "{{ is_state('binary_sensor.test_availability', 'on') }}",
+            },
+        )
+    ],
+)
+@pytest.mark.parametrize(
+    "style", [ConfigurationStyle.MODERN, ConfigurationStyle.TRIGGER]
+)
+@pytest.mark.usefixtures("setup_select")
+async def test_availability(hass: HomeAssistant) -> None:
+    """Test configuration with optimistic state."""
+
+    hass.states.async_set(TEST_AVAILABILITY_ENTITY_ID, "on")
+    hass.states.async_set(TEST_STATE_ENTITY_ID, "test")
+    await hass.async_block_till_done()
+
+    state = hass.states.get(_TEST_SELECT)
+    assert state.state == "test"
+
+    hass.states.async_set(TEST_AVAILABILITY_ENTITY_ID, "off")
+    await hass.async_block_till_done()
+
+    state = hass.states.get(_TEST_SELECT)
+    assert state.state == STATE_UNAVAILABLE
+
+    hass.states.async_set(TEST_STATE_ENTITY_ID, "yes")
+    await hass.async_block_till_done()
+
+    state = hass.states.get(_TEST_SELECT)
+    assert state.state == STATE_UNAVAILABLE
+
+    hass.states.async_set(TEST_AVAILABILITY_ENTITY_ID, "on")
+    await hass.async_block_till_done()
+
+    state = hass.states.get(_TEST_SELECT)
+    assert state.state == "yes"