Restore state of trigger based template binary sensor (#67538)

homeassistant/components/template/binary_sensor.py

@@ -1,7 +1,8 @@
 """Support for exposing a templated binary sensor."""
 from __future__ import annotations
 
-from datetime import timedelta
+from dataclasses import dataclass
+from datetime import datetime, timedelta
 from functools import partial
 import logging
 from typing import Any
@@ -40,9 +41,10 @@ from homeassistant.helpers import template
 import homeassistant.helpers.config_validation as cv
 from homeassistant.helpers.entity import async_generate_entity_id
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
-from homeassistant.helpers.event import async_call_later
+from homeassistant.helpers.event import async_call_later, async_track_point_in_utc_time
-from homeassistant.helpers.restore_state import RestoreEntity
+from homeassistant.helpers.restore_state import ExtraStoredData, RestoreEntity
 from homeassistant.helpers.typing import ConfigType, DiscoveryInfoType
+from homeassistant.util import dt as dt_util
 
 from . import TriggerUpdateCoordinator
 from .const import (
@@ -291,7 +293,7 @@ class BinarySensorTemplate(TemplateEntity, BinarySensorEntity, RestoreEntity):
         return self._device_class
 
 
-class TriggerBinarySensorEntity(TriggerEntity, BinarySensorEntity):
+class TriggerBinarySensorEntity(TriggerEntity, BinarySensorEntity, RestoreEntity):
     """Sensor entity based on trigger data."""
 
     domain = BINARY_SENSOR_DOMAIN
@@ -312,9 +314,36 @@ class TriggerBinarySensorEntity(TriggerEntity, BinarySensorEntity):
                 self._parse_result.add(key)
 
         self._delay_cancel: CALLBACK_TYPE | None = None
-        self._auto_off_cancel = None
+        self._auto_off_cancel: CALLBACK_TYPE | None = None
+        self._auto_off_time: datetime | None = None
         self._state: bool | None = None
 
+    async def async_added_to_hass(self) -> None:
+        """Restore last state."""
+        await super().async_added_to_hass()
+        if (
+            (last_state := await self.async_get_last_state()) is not None
+            and (extra_data := await self.async_get_last_binary_sensor_data())
+            is not None
+            and last_state.state not in (STATE_UNKNOWN, STATE_UNAVAILABLE)
+            # The trigger might have fired already while we waited for stored data,
+            # then we should not restore state
+            and self._state is None
+        ):
+            self._state = last_state.state == STATE_ON
+
+            if CONF_AUTO_OFF not in self._config:
+                return
+
+            if (
+                auto_off_time := extra_data.auto_off_time
+            ) is not None and auto_off_time <= dt_util.utcnow():
+                # It's already past the saved auto off time
+                self._state = False
+
+            if self._state and auto_off_time is not None:
+                self._set_auto_off(auto_off_time)
+
     @property
     def is_on(self) -> bool | None:
         """Return state of the sensor."""
@@ -332,6 +361,7 @@ class TriggerBinarySensorEntity(TriggerEntity, BinarySensorEntity):
         if self._auto_off_cancel:
             self._auto_off_cancel()
             self._auto_off_cancel = None
+            self._auto_off_time = None
 
         if not self.available:
             self.async_write_ha_state()
@@ -372,28 +402,85 @@ class TriggerBinarySensorEntity(TriggerEntity, BinarySensorEntity):
         if not state:
             return
 
-        auto_off_time = self._rendered.get(CONF_AUTO_OFF) or self._config.get(
+        auto_off_delay = self._rendered.get(CONF_AUTO_OFF) or self._config.get(
             CONF_AUTO_OFF
         )
 
-        if auto_off_time is None:
+        if auto_off_delay is None:
             return
 
-        if not isinstance(auto_off_time, timedelta):
+        if not isinstance(auto_off_delay, timedelta):
             try:
-                auto_off_time = cv.positive_time_period(auto_off_time)
+                auto_off_delay = cv.positive_time_period(auto_off_delay)
             except vol.Invalid as err:
                 logging.getLogger(__name__).warning(
                     "Error rendering %s template: %s", CONF_AUTO_OFF, err
                 )
                 return
 
+        auto_off_time = dt_util.utcnow() + auto_off_delay
+        self._set_auto_off(auto_off_time)
+
+    def _set_auto_off(self, auto_off_time: datetime) -> None:
         @callback
         def _auto_off(_):
-            """Set state of template binary sensor."""
+            """Reset state of template binary sensor."""
             self._state = False
             self.async_write_ha_state()
 
-        self._auto_off_cancel = async_call_later(
+        self._auto_off_time = auto_off_time
-            self.hass, auto_off_time.total_seconds(), _auto_off
+        self._auto_off_cancel = async_track_point_in_utc_time(
+            self.hass, _auto_off, self._auto_off_time
         )
+
+    @property
+    def extra_restore_state_data(self) -> AutoOffExtraStoredData:
+        """Return specific state data to be restored."""
+        return AutoOffExtraStoredData(self._auto_off_time)
+
+    async def async_get_last_binary_sensor_data(
+        self,
+    ) -> AutoOffExtraStoredData | None:
+        """Restore auto_off_time."""
+        if (restored_last_extra_data := await self.async_get_last_extra_data()) is None:
+            return None
+        return AutoOffExtraStoredData.from_dict(restored_last_extra_data.as_dict())
+
+
+@dataclass
+class AutoOffExtraStoredData(ExtraStoredData):
+    """Object to hold extra stored data."""
+
+    auto_off_time: datetime | None
+
+    def as_dict(self) -> dict[str, Any]:
+        """Return a dict representation of additional data."""
+        auto_off_time: datetime | None | dict[str, str] = self.auto_off_time
+        if isinstance(auto_off_time, datetime):
+            auto_off_time = {
+                "__type": str(type(auto_off_time)),
+                "isoformat": auto_off_time.isoformat(),
+            }
+        return {
+            "auto_off_time": auto_off_time,
+        }
+
+    @classmethod
+    def from_dict(cls, restored: dict[str, Any]) -> AutoOffExtraStoredData | None:
+        """Initialize a stored binary sensor state from a dict."""
+        try:
+            auto_off_time = restored["auto_off_time"]
+        except KeyError:
+            return None
+        try:
+            type_ = auto_off_time["__type"]
+            if type_ == "<class 'datetime.datetime'>":
+                auto_off_time = dt_util.parse_datetime(auto_off_time["isoformat"])
+        except TypeError:
+            # native_value is not a dict
+            pass
+        except KeyError:
+            # native_value is a dict, but does not have all values
+            return None
+
+        return cls(auto_off_time)

tests/components/template/test_binary_sensor.py

@@ -1,5 +1,5 @@
 """The tests for the Template Binary sensor platform."""
-from datetime import timedelta
+from datetime import datetime, timedelta, timezone
 import logging
 from unittest.mock import patch
 
@@ -24,6 +24,7 @@ from tests.common import (
     assert_setup_component,
     async_fire_time_changed,
     mock_restore_cache,
+    mock_restore_cache_with_extra_data,
 )
 
 ON = "on"
@@ -1112,6 +1113,11 @@ async def test_template_with_trigger_templated_delay_on(hass, start_ha):
     hass.bus.async_fire("test_event", {"beer": 2}, context=context)
     await hass.async_block_till_done()
 
+    # State should still be unknown
+    state = hass.states.get("binary_sensor.test")
+    assert state.state == STATE_UNKNOWN
+
+    # Now wait for the on delay
     future = dt_util.utcnow() + timedelta(seconds=3)
     async_fire_time_changed(hass, future)
     await hass.async_block_till_done()
@@ -1126,3 +1132,170 @@ async def test_template_with_trigger_templated_delay_on(hass, start_ha):
 
     state = hass.states.get("binary_sensor.test")
     assert state.state == OFF
+
+
+@pytest.mark.parametrize("count,domain", [(1, "template")])
+@pytest.mark.parametrize(
+    "config",
+    [
+        {
+            "template": {
+                "trigger": {"platform": "event", "event_type": "test_event"},
+                "binary_sensor": {
+                    "name": "test",
+                    "state": "{{ trigger.event.data.beer == 2 }}",
+                    "device_class": "motion",
+                },
+            },
+        },
+    ],
+)
+@pytest.mark.parametrize(
+    "restored_state, initial_state",
+    [
+        (ON, ON),
+        (OFF, OFF),
+        (STATE_UNAVAILABLE, STATE_UNKNOWN),
+        (STATE_UNKNOWN, STATE_UNKNOWN),
+    ],
+)
+async def test_trigger_entity_restore_state(
+    hass, count, domain, config, restored_state, initial_state
+):
+    """Test restoring trigger template binary sensor."""
+
+    fake_state = State(
+        "binary_sensor.test",
+        restored_state,
+        {},
+    )
+    fake_extra_data = {
+        "auto_off_time": None,
+    }
+    mock_restore_cache_with_extra_data(hass, ((fake_state, fake_extra_data),))
+    with assert_setup_component(count, domain):
+        assert await async_setup_component(
+            hass,
+            domain,
+            config,
+        )
+
+        await hass.async_block_till_done()
+        await hass.async_start()
+        await hass.async_block_till_done()
+
+    state = hass.states.get("binary_sensor.test")
+    assert state.state == initial_state
+
+
+@pytest.mark.parametrize("count,domain", [(1, "template")])
+@pytest.mark.parametrize(
+    "config",
+    [
+        {
+            "template": {
+                "trigger": {"platform": "event", "event_type": "test_event"},
+                "binary_sensor": {
+                    "name": "test",
+                    "state": "{{ trigger.event.data.beer == 2 }}",
+                    "device_class": "motion",
+                    "auto_off": '{{ ({ "seconds": 1 + 1 }) }}',
+                },
+            },
+        },
+    ],
+)
+@pytest.mark.parametrize("restored_state", [ON, OFF])
+async def test_trigger_entity_restore_state_auto_off(
+    hass, count, domain, config, restored_state, freezer
+):
+    """Test restoring trigger template binary sensor."""
+
+    freezer.move_to("2022-02-02 12:02:00+00:00")
+    fake_state = State(
+        "binary_sensor.test",
+        restored_state,
+        {},
+    )
+    fake_extra_data = {
+        "auto_off_time": {
+            "__type": "<class 'datetime.datetime'>",
+            "isoformat": datetime(
+                2022, 2, 2, 12, 2, 2, tzinfo=timezone.utc
+            ).isoformat(),
+        },
+    }
+    mock_restore_cache_with_extra_data(hass, ((fake_state, fake_extra_data),))
+    with assert_setup_component(count, domain):
+        assert await async_setup_component(
+            hass,
+            domain,
+            config,
+        )
+
+        await hass.async_block_till_done()
+        await hass.async_start()
+        await hass.async_block_till_done()
+
+    state = hass.states.get("binary_sensor.test")
+    assert state.state == restored_state
+
+    # Now wait for the auto-off
+    freezer.move_to("2022-02-02 12:02:03+00:00")
+    await hass.async_block_till_done()
+    await hass.async_block_till_done()
+
+    state = hass.states.get("binary_sensor.test")
+    assert state.state == OFF
+
+
+@pytest.mark.parametrize("count,domain", [(1, "template")])
+@pytest.mark.parametrize(
+    "config",
+    [
+        {
+            "template": {
+                "trigger": {"platform": "event", "event_type": "test_event"},
+                "binary_sensor": {
+                    "name": "test",
+                    "state": "{{ trigger.event.data.beer == 2 }}",
+                    "device_class": "motion",
+                    "auto_off": '{{ ({ "seconds": 1 + 1 }) }}',
+                },
+            },
+        },
+    ],
+)
+async def test_trigger_entity_restore_state_auto_off_expired(
+    hass, count, domain, config, freezer
+):
+    """Test restoring trigger template binary sensor."""
+
+    freezer.move_to("2022-02-02 12:02:00+00:00")
+    fake_state = State(
+        "binary_sensor.test",
+        ON,
+        {},
+    )
+    fake_extra_data = {
+        "auto_off_time": {
+            "__type": "<class 'datetime.datetime'>",
+            "isoformat": datetime(
+                2022, 2, 2, 12, 2, 0, tzinfo=timezone.utc
+            ).isoformat(),
+        },
+    }
+    mock_restore_cache_with_extra_data(hass, ((fake_state, fake_extra_data),))
+    with assert_setup_component(count, domain):
+        assert await async_setup_component(
+            hass,
+            domain,
+            config,
+        )
+
+        await hass.async_block_till_done()
+        await hass.async_start()
+        await hass.async_block_till_done()
+
+    state = hass.states.get("binary_sensor.test")
+    assert state.state == OFF