From 7fc0ffd5c591429cef805aca707acdda0ca304e6 Mon Sep 17 00:00:00 2001 From: Erik Montnemery Date: Mon, 14 Mar 2022 20:28:55 +0100 Subject: [PATCH] Restore state of trigger based template binary sensor (#67538) --- .../components/template/binary_sensor.py | 111 +++++++++-- .../components/template/test_binary_sensor.py | 175 +++++++++++++++++- 2 files changed, 273 insertions(+), 13 deletions(-) diff --git a/homeassistant/components/template/binary_sensor.py b/homeassistant/components/template/binary_sensor.py index f416454f388..53267b29341 100644 --- a/homeassistant/components/template/binary_sensor.py +++ b/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.restore_state import RestoreEntity +from homeassistant.helpers.event import async_call_later, async_track_point_in_utc_time +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.hass, auto_off_time.total_seconds(), _auto_off + self._auto_off_time = auto_off_time + 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_ == "": + 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) diff --git a/tests/components/template/test_binary_sensor.py b/tests/components/template/test_binary_sensor.py index c3773e8eb13..1d8d9473f3c 100644 --- a/tests/components/template/test_binary_sensor.py +++ b/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": "", + "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": "", + "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