Allow defining state class for template sensors (#52130)

This commit is contained in:
Paulus Schoutsen 2021-06-23 14:37:04 -07:00 committed by GitHub
parent 6352d8fb0e
commit cc00617cd5
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
6 changed files with 70 additions and 73 deletions

View File

@ -25,3 +25,4 @@ CONF_AVAILABILITY = "availability"
CONF_ATTRIBUTES = "attributes"
CONF_PICTURE = "picture"
CONF_OBJECT_ID = "object_id"
CONF_STATE_CLASS = "state_class"

View File

@ -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):

View File

@ -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)

View File

@ -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()

View File

@ -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

View File

@ -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"