From 9bec6493230f6c64d51ba7f45edc7ff94999e5fd Mon Sep 17 00:00:00 2001 From: Erik Montnemery Date: Thu, 21 Apr 2022 18:32:18 +0200 Subject: [PATCH] Restore state of trigger-based template sensor (#69344) --- homeassistant/components/template/sensor.py | 19 +++- tests/components/template/test_sensor.py | 97 ++++++++++++++++++++- 2 files changed, 113 insertions(+), 3 deletions(-) diff --git a/homeassistant/components/template/sensor.py b/homeassistant/components/template/sensor.py index a61d15cb7a4..126dd551c45 100644 --- a/homeassistant/components/template/sensor.py +++ b/homeassistant/components/template/sensor.py @@ -13,6 +13,7 @@ from homeassistant.components.sensor import ( ENTITY_ID_FORMAT, PLATFORM_SCHEMA, STATE_CLASSES_SCHEMA, + RestoreSensor, SensorDeviceClass, SensorEntity, ) @@ -30,6 +31,8 @@ from homeassistant.const import ( CONF_UNIQUE_ID, CONF_UNIT_OF_MEASUREMENT, CONF_VALUE_TEMPLATE, + STATE_UNAVAILABLE, + STATE_UNKNOWN, ) from homeassistant.core import HomeAssistant, callback from homeassistant.exceptions import TemplateError @@ -237,12 +240,26 @@ class SensorTemplate(TemplateEntity, SensorEntity): ) -class TriggerSensorEntity(TriggerEntity, SensorEntity): +class TriggerSensorEntity(TriggerEntity, RestoreSensor): """Sensor entity based on trigger data.""" domain = SENSOR_DOMAIN extra_template_keys = (CONF_STATE,) + 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_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 CONF_STATE not in self._rendered + ): + self._rendered[CONF_STATE] = extra_data.native_value + self.restore_attributes(last_state) + @property def native_value(self) -> str | datetime | date | None: """Return state of the sensor.""" diff --git a/tests/components/template/test_sensor.py b/tests/components/template/test_sensor.py index 033d5636720..297008e77bb 100644 --- a/tests/components/template/test_sensor.py +++ b/tests/components/template/test_sensor.py @@ -17,14 +17,18 @@ from homeassistant.const import ( STATE_UNAVAILABLE, STATE_UNKNOWN, ) -from homeassistant.core import Context, CoreState, callback +from homeassistant.core import Context, CoreState, State, callback from homeassistant.helpers import entity_registry from homeassistant.helpers.entity_component import async_update_entity from homeassistant.helpers.template import Template from homeassistant.setup import ATTR_COMPONENT, async_setup_component import homeassistant.util.dt as dt_util -from tests.common import async_fire_time_changed +from tests.common import ( + assert_setup_component, + async_fire_time_changed, + mock_restore_cache_with_extra_data, +) TEST_NAME = "sensor.test_template_sensor" @@ -1288,3 +1292,92 @@ async def test_entity_device_class_errors_works(hass): ts_state = hass.states.get("sensor.timestamp_entity") assert ts_state is not None assert ts_state.state == STATE_UNKNOWN + + +@pytest.mark.parametrize("count,domain", [(1, "template")]) +@pytest.mark.parametrize( + "config", + [ + { + "template": { + "trigger": {"platform": "event", "event_type": "test_event"}, + "sensor": { + "name": "test", + "state": "{{ trigger.event.data.beer }}", + "picture": "{{ '/local/dogs.png' }}", + "icon": "{{ 'mdi:pirate' }}", + "attributes": { + "plus_one": "{{ trigger.event.data.beer + 1 }}", + "another": "{{ trigger.event.data.uno_mas or 1 }}", + }, + }, + }, + }, + ], +) +@pytest.mark.parametrize( + "restored_state, restored_native_value, initial_state, initial_attributes", + [ + # the native value should be used, not the state + ("dog", 10, "10", ["entity_picture", "icon", "plus_one"]), + (STATE_UNAVAILABLE, 10, STATE_UNKNOWN, []), + (STATE_UNKNOWN, 10, STATE_UNKNOWN, []), + ], +) +async def test_trigger_entity_restore_state( + hass, + count, + domain, + config, + restored_state, + restored_native_value, + initial_state, + initial_attributes, +): + """Test restoring trigger template binary sensor.""" + + restored_attributes = { + "entity_picture": "/local/cats.png", + "icon": "mdi:ship", + "plus_one": 55, + } + + fake_state = State( + "sensor.test", + restored_state, + restored_attributes, + ) + fake_extra_data = { + "native_value": restored_native_value, + "native_unit_of_measurement": 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("sensor.test") + assert state.state == initial_state + for attr in restored_attributes: + if attr in initial_attributes: + assert state.attributes[attr] == restored_attributes[attr] + else: + assert attr not in state.attributes + assert "another" not in state.attributes + + hass.bus.async_fire("test_event", {"beer": 2}) + await hass.async_block_till_done() + + state = hass.states.get("sensor.test") + assert state.state == "2" + assert state.attributes["icon"] == "mdi:pirate" + assert state.attributes["entity_picture"] == "/local/dogs.png" + assert state.attributes["plus_one"] == 3 + assert state.attributes["another"] == 1