From d94e20b2f1b396e0caeda63b2b4554ceeb041de2 Mon Sep 17 00:00:00 2001 From: Petro31 <35082313+Petro31@users.noreply.github.com> Date: Wed, 29 Jul 2026 05:58:42 -0400 Subject: [PATCH] Add restore to lock template entities (#176571) --- homeassistant/components/template/lock.py | 67 ++++++++- tests/components/template/test_lock.py | 157 ++++++++++++++++++++++ 2 files changed, 222 insertions(+), 2 deletions(-) diff --git a/homeassistant/components/template/lock.py b/homeassistant/components/template/lock.py index a77f76aef9c0..9c31be9cba50 100644 --- a/homeassistant/components/template/lock.py +++ b/homeassistant/components/template/lock.py @@ -1,6 +1,7 @@ """Support for locks which integrates with other components.""" -from typing import TYPE_CHECKING, Any, override +from dataclasses import asdict, dataclass +from typing import TYPE_CHECKING, Any, Self, override import voluptuous as vol @@ -20,6 +21,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 validators as template_validators @@ -122,12 +124,48 @@ def async_create_preview_lock( ) -class AbstractTemplateLock(AbstractTemplateEntity, LockEntity): +@dataclass(kw_only=True) +class LockExtraStoredData(ExtraStoredData): + """Holds extra stored data for template lock entities.""" + + code_format: str | None + is_locked: bool | None + is_locking: bool | None + is_open: bool | None + is_opening: bool | None + is_unlocking: bool | None + is_jammed: bool | None + + @override + def as_dict(self) -> dict[str, Any]: + """Return a dict representation of the lock data.""" + return asdict(self) + + @classmethod + def from_dict(cls, restored: dict[str, Any]) -> Self | None: + """Initialize a stored lock state from a dict.""" + try: + return cls( + code_format=restored["code_format"], + is_locked=restored["is_locked"], + is_locking=restored["is_locking"], + is_open=restored["is_open"], + is_opening=restored["is_opening"], + is_unlocking=restored["is_unlocking"], + is_jammed=restored["is_jammed"], + ) + except KeyError: + return None + + +class AbstractTemplateLock(AbstractTemplateEntity, LockEntity, RestoreEntity): """Representation of a template lock features.""" _entity_id_format = ENTITY_ID_FORMAT _optimistic_entity = True _state_option = CONF_STATE + _restore_state_extra_data = LockExtraStoredData + _restore_state_properties = ("_attr_is_locked",) # The super init is not called because TemplateEntity # and TriggerEntity will call @@ -260,6 +298,31 @@ class AbstractTemplateLock(AbstractTemplateEntity, LockEntity): }, ) + @property + @override + def extra_restore_state_data(self) -> LockExtraStoredData: + """Return lock specific state data to be restored.""" + return LockExtraStoredData( + code_format=self._attr_code_format, + is_locked=self._attr_is_locked, + is_locking=self._attr_is_locking, + is_open=self._attr_is_open, + is_opening=self._attr_is_opening, + is_unlocking=self._attr_is_unlocking, + is_jammed=self._attr_is_jammed, + ) + + @override + def restore_extra_data(self, extra_data: LockExtraStoredData) -> None: + """Restore the extra data.""" + self._attr_code_format = extra_data.code_format + self._attr_is_locked = extra_data.is_locked + self._attr_is_locking = extra_data.is_locking + self._attr_is_open = extra_data.is_open + self._attr_is_opening = extra_data.is_opening + self._attr_is_unlocking = extra_data.is_unlocking + self._attr_is_jammed = extra_data.is_jammed + class StateLockEntity(TemplateEntity, AbstractTemplateLock): """Representation of a template lock.""" diff --git a/tests/components/template/test_lock.py b/tests/components/template/test_lock.py index 85bd86cc46b6..07443e251043 100644 --- a/tests/components/template/test_lock.py +++ b/tests/components/template/test_lock.py @@ -24,6 +24,7 @@ from homeassistant.helpers.typing import ConfigType from .conftest import ( ConfigurationStyle, TemplatePlatformSetup, + assert_state_and_attributes, async_get_flow_preview_state, async_trigger, make_test_action, @@ -31,6 +32,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 @@ -998,3 +1001,157 @@ async def test_flow_preview( ) assert state["state"] == LockState.LOCKED + + +@pytest.mark.parametrize( + "style", [ConfigurationStyle.MODERN, ConfigurationStyle.TRIGGER] +) +@pytest.mark.parametrize( + ( + "saved_state", + "saved_extra_data", + "initial_state", + "initial_attributes", + ), + [ + ( + LockState.JAMMED, + { + "code_format": ".+", + "is_locked": False, + "is_locking": False, + "is_open": False, + "is_opening": False, + "is_unlocking": False, + "is_jammed": True, + }, + LockState.JAMMED, + { + "code_format": ".+", + }, + ), + ( + LockState.LOCKED, + { + "code_format": ".+", + "is_locked": True, + "is_locking": False, + "is_open": False, + "is_opening": False, + "is_unlocking": False, + "is_jammed": False, + }, + LockState.LOCKED, + { + "code_format": ".+", + }, + ), + ( + LockState.LOCKING, + { + "code_format": ".+", + "is_locked": False, + "is_locking": True, + "is_open": False, + "is_opening": False, + "is_unlocking": False, + "is_jammed": False, + }, + LockState.LOCKING, + { + "code_format": ".+", + }, + ), + ( + LockState.JAMMED, + { + "code_format": ".+", + "is_locked": False, + "is_locking": False, + "is_opening": False, + "is_unlocking": False, + "is_jammed": True, + }, + STATE_UNKNOWN, + { + "code_format": None, + }, + ), + ( + STATE_UNAVAILABLE, + { + "code_format": ".+", + "is_locked": False, + "is_locking": False, + "is_open": False, + "is_opening": False, + "is_unlocking": False, + "is_jammed": True, + }, + STATE_UNKNOWN, + { + "code_format": None, + }, + ), + ( + STATE_UNKNOWN, + { + "code_format": ".+", + "is_locked": False, + "is_locking": False, + "is_open": False, + "is_opening": False, + "is_unlocking": False, + "is_jammed": True, + }, + STATE_UNKNOWN, + { + "code_format": None, + }, + ), + ], +) +async def test_restore_state( + hass: HomeAssistant, + style: ConfigurationStyle, + saved_state: LockState | str, + saved_extra_data: dict | None, + initial_state: LockState | str, + initial_attributes: ConfigType, +) -> None: + """Test restoring state.""" + + setup_mock_template_entity_restore_state( + hass, + TEST_LOCK, + saved_state, + saved_extra_data=saved_extra_data, + ) + + await setup_restore_template_entity( + hass, + TEST_LOCK, + style, + { + "code_format": "{{ state_attr('sensor.test_state', 'code_format') }}", + "state": "{{ state_attr('sensor.test_state', 'lock_state') }}", + "lock": [], + "open": [], + "unlock": [], + }, + "is_state_attr('sensor.test_state', 'lock_state', 'unlocked')", + ) + + assert_state_and_attributes(hass, TEST_LOCK, initial_state, initial_attributes) + + await async_trigger( + hass, + "sensor.test_state", + "anything", + {"lock_state": LockState.UNLOCKED, "code_format": "\\\\d+"}, + ) + + # The first trigger should replace the restored code_format attribute + assert_state_and_attributes( + hass, TEST_LOCK, LockState.UNLOCKED, {"code_format": "\\\\d+"} + )