Add restore to select template entities (#176692)

homeassistant/components/template/select.py

@@ -1,7 +1,8 @@
 """Support for selects which integrates with other components."""
 
+from dataclasses import asdict, dataclass
 import logging
-from typing import TYPE_CHECKING, Any, override
+from typing import TYPE_CHECKING, Any, Self, override
 
 import voluptuous as vol
 
@@ -18,6 +19,7 @@ from homeassistant.helpers.entity_platform import (
     AddConfigEntryEntitiesCallback,
     AddEntitiesCallback,
 )
+from homeassistant.helpers.restore_state import ExtraStoredData, RestoreEntity
 from homeassistant.helpers.typing import ConfigType, DiscoveryInfoType
 
 from . import TriggerUpdateCoordinator, validators as template_validators
@@ -106,12 +108,38 @@ def async_create_preview_select(
     )
 
 
-class AbstractTemplateSelect(AbstractTemplateEntity, SelectEntity):
+@dataclass(kw_only=True)
+class SelectExtraStoredData(ExtraStoredData):
+    """Holds extra stored data for template select entities."""
+
+    current_option: str | None
+    options: list[str]
+
+    @override
+    def as_dict(self) -> dict[str, Any]:
+        """Return a dict representation of the select data."""
+        return asdict(self)
+
+    @classmethod
+    def from_dict(cls, restored: dict[str, Any]) -> Self | None:
+        """Initialize a stored select state from a dict."""
+        try:
+            return cls(
+                current_option=restored["current_option"],
+                options=restored["options"],
+            )
+        except KeyError:
+            return None
+
+
+class AbstractTemplateSelect(AbstractTemplateEntity, SelectEntity, RestoreEntity):
     """Representation of a template select features."""
 
     _entity_id_format = ENTITY_ID_FORMAT
     _optimistic_entity = True
     _state_option = CONF_STATE
+    _restore_state_extra_data = SelectExtraStoredData
+    _restore_state_properties = ("_attr_current_option",)
 
     # The super init is not called because TemplateEntity
     # and TriggerEntity will call
@@ -124,7 +152,7 @@ class AbstractTemplateSelect(AbstractTemplateEntity, SelectEntity):
 
         self.setup_state_template(
             "_attr_current_option",
-            cv.string,
+            template_validators.string(self, CONF_STATE),
         )
         self.setup_template(
             CONF_OPTIONS,
@@ -150,6 +178,21 @@ class AbstractTemplateSelect(AbstractTemplateEntity, SelectEntity):
                 context=self._context,
             )
 
+    @property
+    @override
+    def extra_restore_state_data(self) -> SelectExtraStoredData:
+        """Return select specific state data to be restored."""
+        return SelectExtraStoredData(
+            current_option=self._attr_current_option,
+            options=self._attr_options or [],
+        )
+
+    @override
+    def restore_extra_data(self, extra_data: SelectExtraStoredData) -> None:
+        """Restore the extra data."""
+        self._attr_current_option = extra_data.current_option
+        self._attr_options = extra_data.options
+
 
 class TemplateSelect(TemplateEntity, AbstractTemplateSelect):
     """Representation of a template select."""

tests/components/template/test_select.py

@@ -26,11 +26,13 @@ from homeassistant.const import (
 )
 from homeassistant.core import HomeAssistant, ServiceCall
 from homeassistant.helpers import device_registry as dr, entity_registry as er
+from homeassistant.helpers.typing import ConfigType
 
 from .conftest import (
     ConfigurationStyle,
     TemplatePlatformSetup,
     assert_action,
+    assert_state_and_attributes,
     async_get_flow_preview_state,
     async_trigger,
     make_test_action,
@@ -38,6 +40,8 @@ from .conftest import (
     setup_and_test_nested_unique_id,
     setup_and_test_unique_id,
     setup_entity,
+    setup_mock_template_entity_restore_state,
+    setup_restore_template_entity,
 )
 
 from tests.common import MockConfigEntry, assert_setup_component
@@ -213,6 +217,9 @@ async def test_template_select(hass: HomeAssistant, calls: list[ServiceCall]) ->
     await async_trigger(hass, TEST_STATE_ENTITY_ID, "c", attributes)
     _verify(hass, "c", ["a", "b", "c"])
 
+    await async_trigger(hass, TEST_STATE_ENTITY_ID, "None", attributes)
+    _verify(hass, STATE_UNKNOWN, ["a", "b", "c"])
+
 
 def _verify(
     hass: HomeAssistant,
@@ -556,3 +563,125 @@ async def test_nested_unique_id(
         TEST_OPTIONS_WITHOUT_STATE,
         "{{ 'test' }}",
     )
+
+
+@pytest.mark.parametrize(
+    "style", [ConfigurationStyle.MODERN, ConfigurationStyle.TRIGGER]
+)
+@pytest.mark.parametrize(
+    (
+        "saved_state",
+        "saved_extra_data",
+        "initial_state",
+        "initial_attributes",
+    ),
+    [
+        (
+            "something",
+            {
+                "current_option": "something",
+                "options": ["something", "anything"],
+            },
+            "something",
+            {
+                "options": ["something", "anything"],
+            },
+        ),
+        (
+            "something",
+            {
+                "current_option": "something",
+            },
+            STATE_UNKNOWN,
+            {
+                "options": [],
+            },
+        ),
+        (
+            "something",
+            {
+                "options": ["something", "anything"],
+            },
+            STATE_UNKNOWN,
+            {
+                "options": [],
+            },
+        ),
+        (
+            STATE_UNAVAILABLE,
+            {
+                "current_option": "something",
+                "options": ["something", "anything"],
+            },
+            STATE_UNKNOWN,
+            {
+                "options": [],
+            },
+        ),
+        (
+            STATE_UNKNOWN,
+            {
+                "current_option": "something",
+                "options": ["something", "anything"],
+            },
+            STATE_UNKNOWN,
+            {
+                "options": [],
+            },
+        ),
+    ],
+)
+async def test_restore_state(
+    hass: HomeAssistant,
+    style: ConfigurationStyle,
+    saved_state: str,
+    saved_extra_data: dict | None,
+    initial_state: str,
+    initial_attributes: ConfigType,
+) -> None:
+    """Test restoring state."""
+
+    setup_mock_template_entity_restore_state(
+        hass,
+        TEST_SELECT,
+        saved_state,
+        saved_extra_data=saved_extra_data,
+    )
+
+    await setup_restore_template_entity(
+        hass,
+        TEST_SELECT,
+        style,
+        {
+            "state": "{{ state_attr('sensor.test_state', 'option') }}",
+            "options": "{{ state_attr('sensor.test_state', 'options') or [] }}",
+            "select_option": [],
+        },
+        "is_state('sensor.test_state', 'something_new')",
+    )
+
+    assert_state_and_attributes(
+        hass,
+        TEST_SELECT,
+        initial_state,
+        initial_attributes,
+    )
+
+    await async_trigger(
+        hass,
+        "sensor.test_state",
+        "anything",
+        {
+            "options": ["something", "anything", "something_new"],
+            "option": "something_new",
+        },
+    )
+
+    assert_state_and_attributes(
+        hass,
+        TEST_SELECT,
+        "something_new",
+        {
+            "options": ["something", "anything", "something_new"],
+        },
+    )