From cc00617cd52a9c453676b9c7978115295f181fd8 Mon Sep 17 00:00:00 2001 From: Paulus Schoutsen Date: Wed, 23 Jun 2021 14:37:04 -0700 Subject: [PATCH] Allow defining state class for template sensors (#52130) --- homeassistant/components/template/const.py | 1 + homeassistant/components/template/light.py | 2 +- homeassistant/components/template/sensor.py | 58 ++++++++----------- .../components/template/template_entity.py | 50 +++++----------- homeassistant/components/template/vacuum.py | 2 +- tests/components/template/test_sensor.py | 30 ++++++++++ 6 files changed, 70 insertions(+), 73 deletions(-) diff --git a/homeassistant/components/template/const.py b/homeassistant/components/template/const.py index 661953bcfa5..31896e930e4 100644 --- a/homeassistant/components/template/const.py +++ b/homeassistant/components/template/const.py @@ -25,3 +25,4 @@ CONF_AVAILABILITY = "availability" CONF_ATTRIBUTES = "attributes" CONF_PICTURE = "picture" CONF_OBJECT_ID = "object_id" +CONF_STATE_CLASS = "state_class" diff --git a/homeassistant/components/template/light.py b/homeassistant/components/template/light.py index f546c5dc4da..b8ebe03ceba 100644 --- a/homeassistant/components/template/light.py +++ b/homeassistant/components/template/light.py @@ -595,7 +595,7 @@ class LightTemplate(TemplateEntity, LightEntity): # This behavior is legacy self._state = False if not self._availability_template: - self._available = True + self._attr_available = True return if isinstance(result, bool): diff --git a/homeassistant/components/template/sensor.py b/homeassistant/components/template/sensor.py index 56e0e11edb0..a887890510a 100644 --- a/homeassistant/components/template/sensor.py +++ b/homeassistant/components/template/sensor.py @@ -8,6 +8,7 @@ from homeassistant.components.sensor import ( DOMAIN as SENSOR_DOMAIN, ENTITY_ID_FORMAT, PLATFORM_SCHEMA, + STATE_CLASSES_SCHEMA, SensorEntity, ) from homeassistant.const import ( @@ -37,6 +38,7 @@ from .const import ( CONF_AVAILABILITY_TEMPLATE, CONF_OBJECT_ID, CONF_PICTURE, + CONF_STATE_CLASS, CONF_TRIGGER, ) from .template_entity import TemplateEntity @@ -64,6 +66,7 @@ SENSOR_SCHEMA = vol.Schema( vol.Optional(CONF_UNIT_OF_MEASUREMENT): cv.string, vol.Optional(CONF_DEVICE_CLASS): DEVICE_CLASSES_SCHEMA, vol.Optional(CONF_UNIQUE_ID): cv.string, + vol.Optional(CONF_STATE_CLASS): STATE_CLASSES_SCHEMA, } ) @@ -159,6 +162,7 @@ def _async_create_template_tracking_entities( device_class = entity_conf.get(CONF_DEVICE_CLASS) attribute_templates = entity_conf.get(CONF_ATTRIBUTES, {}) unique_id = entity_conf.get(CONF_UNIQUE_ID) + state_class = entity_conf.get(CONF_STATE_CLASS) if unique_id and unique_id_prefix: unique_id = f"{unique_id_prefix}-{unique_id}" @@ -176,6 +180,7 @@ def _async_create_template_tracking_entities( device_class, attribute_templates, unique_id, + state_class, ) ) @@ -224,6 +229,7 @@ class SensorTemplate(TemplateEntity, SensorEntity): device_class: str | None, attribute_templates: dict[str, template.Template], unique_id: str | None, + state_class: str | None, ) -> None: """Initialize the sensor.""" super().__init__( @@ -237,61 +243,38 @@ class SensorTemplate(TemplateEntity, SensorEntity): ENTITY_ID_FORMAT, object_id, hass=hass ) - self._name: str | None = None self._friendly_name_template = friendly_name_template # Try to render the name as it can influence the entity ID if friendly_name_template: friendly_name_template.hass = hass try: - self._name = friendly_name_template.async_render(parse_result=False) + self._attr_name = friendly_name_template.async_render( + parse_result=False + ) except template.TemplateError: pass - self._unit_of_measurement = unit_of_measurement + self._attr_unit_of_measurement = unit_of_measurement self._template = state_template - self._state = None - self._device_class = device_class - - self._unique_id = unique_id + self._attr_device_class = device_class + self._attr_state_class = state_class + self._attr_unique_id = unique_id async def async_added_to_hass(self): """Register callbacks.""" - self.add_template_attribute("_state", self._template, None, self._update_state) + self.add_template_attribute( + "_attr_state", self._template, None, self._update_state + ) if self._friendly_name_template and not self._friendly_name_template.is_static: - self.add_template_attribute("_name", self._friendly_name_template) + self.add_template_attribute("_attr_name", self._friendly_name_template) await super().async_added_to_hass() @callback def _update_state(self, result): super()._update_state(result) - self._state = None if isinstance(result, TemplateError) else result - - @property - def name(self): - """Return the name of the sensor.""" - return self._name - - @property - def unique_id(self): - """Return the unique id of this sensor.""" - return self._unique_id - - @property - def state(self): - """Return the state of the sensor.""" - return self._state - - @property - def device_class(self) -> str | None: - """Return the device class of the sensor.""" - return self._device_class - - @property - def unit_of_measurement(self): - """Return the unit_of_measurement of the device.""" - return self._unit_of_measurement + self._attr_state = None if isinstance(result, TemplateError) else result class TriggerSensorEntity(TriggerEntity, SensorEntity): @@ -304,3 +287,8 @@ class TriggerSensorEntity(TriggerEntity, SensorEntity): def state(self) -> str | None: """Return state of the sensor.""" return self._rendered.get(CONF_STATE) + + @property + def state_class(self) -> str | None: + """Sensor state class.""" + return self._config.get(CONF_STATE_CLASS) diff --git a/homeassistant/components/template/template_entity.py b/homeassistant/components/template/template_entity.py index 522eb7d89ba..7bf6d6109be 100644 --- a/homeassistant/components/template/template_entity.py +++ b/homeassistant/components/template/template_entity.py @@ -112,6 +112,8 @@ class _TemplateAttribute: class TemplateEntity(Entity): """Entity that uses templates to calculate attributes.""" + _attr_should_poll = False + def __init__( self, *, @@ -124,54 +126,27 @@ class TemplateEntity(Entity): self._template_attrs = {} self._async_update = None self._attribute_templates = attribute_templates - self._attributes = {} + self._attr_extra_state_attributes = {} self._availability_template = availability_template - self._available = True + self._attr_available = True self._icon_template = icon_template self._entity_picture_template = entity_picture_template - self._icon = None - self._entity_picture = None self._self_ref_update_count = 0 - @property - def should_poll(self): - """No polling needed.""" - return False - @callback def _update_available(self, result): if isinstance(result, TemplateError): - self._available = True + self._attr_available = True return - self._available = result_as_boolean(result) + self._attr_available = result_as_boolean(result) @callback def _update_state(self, result): if self._availability_template: return - self._available = not isinstance(result, TemplateError) - - @property - def available(self) -> bool: - """Return if the device is available.""" - return self._available - - @property - def icon(self): - """Return the icon to use in the frontend, if any.""" - return self._icon - - @property - def entity_picture(self): - """Return the entity_picture to use in the frontend, if any.""" - return self._entity_picture - - @property - def extra_state_attributes(self): - """Return the state attributes.""" - return self._attributes + self._attr_available = not isinstance(result, TemplateError) @callback def _add_attribute_template(self, attribute_key, attribute_template): @@ -179,7 +154,7 @@ class TemplateEntity(Entity): def _update_attribute(result): attr_result = None if isinstance(result, TemplateError) else result - self._attributes[attribute_key] = attr_result + self._attr_extra_state_attributes[attribute_key] = attr_result self.add_template_attribute( attribute_key, attribute_template, None, _update_attribute @@ -271,18 +246,21 @@ class TemplateEntity(Entity): """Run when entity about to be added to hass.""" if self._availability_template is not None: self.add_template_attribute( - "_available", self._availability_template, None, self._update_available + "_attr_available", + self._availability_template, + None, + self._update_available, ) if self._attribute_templates is not None: for key, value in self._attribute_templates.items(): self._add_attribute_template(key, value) if self._icon_template is not None: self.add_template_attribute( - "_icon", self._icon_template, vol.Or(cv.whitespace, cv.icon) + "_attr_icon", self._icon_template, vol.Or(cv.whitespace, cv.icon) ) if self._entity_picture_template is not None: self.add_template_attribute( - "_entity_picture", self._entity_picture_template + "_attr_entity_picture", self._entity_picture_template ) if self.hass.state == CoreState.running: await self._async_template_startup() diff --git a/homeassistant/components/template/vacuum.py b/homeassistant/components/template/vacuum.py index ed7919d174e..78c51d2009c 100644 --- a/homeassistant/components/template/vacuum.py +++ b/homeassistant/components/template/vacuum.py @@ -362,7 +362,7 @@ class TemplateVacuum(TemplateEntity, StateVacuumEntity): # This is legacy behavior self._state = STATE_UNKNOWN if not self._availability_template: - self._available = True + self._attr_available = True return # Validate state diff --git a/tests/components/template/test_sensor.py b/tests/components/template/test_sensor.py index 4047a822432..df5c43aa58b 100644 --- a/tests/components/template/test_sensor.py +++ b/tests/components/template/test_sensor.py @@ -1044,6 +1044,7 @@ async def test_trigger_entity(hass): "attributes": { "plus_one": "{{ trigger.event.data.beer + 1 }}" }, + "state_class": "measurement", } ], }, @@ -1100,6 +1101,7 @@ async def test_trigger_entity(hass): assert state.attributes.get("entity_picture") == "/local/dogs.png" assert state.attributes.get("plus_one") == 3 assert state.attributes.get("unit_of_measurement") == "%" + assert state.attributes.get("state_class") == "measurement" assert state.context is context @@ -1167,3 +1169,31 @@ async def test_trigger_not_allowed_platform_config(hass, caplog): "You can only add triggers to template entities if they are defined under `template:`." in caplog.text ) + + +async def test_config_top_level(hass): + """Test unique_id option only creates one sensor per id.""" + await async_setup_component( + hass, + "template", + { + "template": { + "sensor": { + "name": "top-level", + "device_class": "battery", + "state_class": "measurement", + "state": "5", + "unit_of_measurement": "%", + }, + }, + }, + ) + + await hass.async_block_till_done() + + assert len(hass.states.async_all()) == 1 + state = hass.states.get("sensor.top_level") + assert state is not None + assert state.state == "5" + assert state.attributes["device_class"] == "battery" + assert state.attributes["state_class"] == "measurement"