New template sensor attributes (home-assistant#26127)

* updated sensor and test files

* Formatting fixes

* Updated attribute template code

* Black formatting

* Code improvements based on feedback on binary_sensor pull request

* Updated tests

* Remove duplicated code and fix tests

* Black formatting on tests

* Remove link from docstring

* Moved default to schema

* Formatting fix and change to use dict[key] to retrieve attribute_templates

homeassistant/components/template/sensor.py

@@ -1,6 +1,7 @@
 """Allows the creation of a sensor that breaks out state_attributes."""
 import logging
 from typing import Optional
+from itertools import chain
 
 import voluptuous as vol
 
@@ -28,6 +29,8 @@
 from homeassistant.helpers.entity import Entity, async_generate_entity_id
 from homeassistant.helpers.event import async_track_state_change
 
+CONF_ATTRIBUTE_TEMPLATES = "attribute_templates"
+
 _LOGGER = logging.getLogger(__name__)
 
 SENSOR_SCHEMA = vol.Schema(
@@ -36,6 +39,9 @@
         vol.Optional(CONF_ICON_TEMPLATE): cv.template,
         vol.Optional(CONF_ENTITY_PICTURE_TEMPLATE): cv.template,
         vol.Optional(CONF_FRIENDLY_NAME_TEMPLATE): cv.template,
+        vol.Optional(CONF_ATTRIBUTE_TEMPLATES, default={}): vol.Schema(
+            {cv.string: cv.template}
+        ),
         vol.Optional(ATTR_FRIENDLY_NAME): cv.string,
         vol.Optional(ATTR_UNIT_OF_MEASUREMENT): cv.string,
         vol.Optional(CONF_DEVICE_CLASS): DEVICE_CLASSES_SCHEMA,
@@ -60,17 +66,20 @@ async def async_setup_platform(hass, config, async_add_entities, discovery_info=
         friendly_name_template = device_config.get(CONF_FRIENDLY_NAME_TEMPLATE)
         unit_of_measurement = device_config.get(ATTR_UNIT_OF_MEASUREMENT)
         device_class = device_config.get(CONF_DEVICE_CLASS)
+        attribute_templates = device_config[CONF_ATTRIBUTE_TEMPLATES]
 
         entity_ids = set()
         manual_entity_ids = device_config.get(ATTR_ENTITY_ID)
         invalid_templates = []
 
-        for tpl_name, template in (
-            (CONF_VALUE_TEMPLATE, state_template),
-            (CONF_ICON_TEMPLATE, icon_template),
-            (CONF_ENTITY_PICTURE_TEMPLATE, entity_picture_template),
-            (CONF_FRIENDLY_NAME_TEMPLATE, friendly_name_template),
-        ):
+        templates = {
+            CONF_VALUE_TEMPLATE: state_template,
+            CONF_ICON_TEMPLATE: icon_template,
+            CONF_ENTITY_PICTURE_TEMPLATE: entity_picture_template,
+            CONF_FRIENDLY_NAME_TEMPLATE: friendly_name_template,
+        }
+
+        for tpl_name, template in chain(templates.items(), attribute_templates.items()):
             if template is None:
                 continue
             template.hass = hass
@@ -82,7 +91,7 @@ async def async_setup_platform(hass, config, async_add_entities, discovery_info=
             if template_entity_ids == MATCH_ALL:
                 entity_ids = MATCH_ALL
                 # Cut off _template from name
-                invalid_templates.append(tpl_name[:-9])
+                invalid_templates.append(tpl_name.replace("_template", ""))
             elif entity_ids != MATCH_ALL:
                 entity_ids |= set(template_entity_ids)
 
@@ -113,6 +122,7 @@ async def async_setup_platform(hass, config, async_add_entities, discovery_info=
                 entity_picture_template,
                 entity_ids,
                 device_class,
+                attribute_templates,
             )
         )
     if not sensors:
@@ -138,6 +148,7 @@ def __init__(
         entity_picture_template,
         entity_ids,
         device_class,
+        attribute_templates,
     ):
         """Initialize the sensor."""
         self.hass = hass
@@ -155,6 +166,8 @@ def __init__(
         self._entity_picture = None
         self._entities = entity_ids
         self._device_class = device_class
+        self._attribute_templates = attribute_templates
+        self._attributes = {}
 
     async def async_added_to_hass(self):
         """Register callbacks."""
@@ -209,6 +222,11 @@ def unit_of_measurement(self):
         """Return the unit_of_measurement of the device."""
         return self._unit_of_measurement
 
+    @property
+    def device_state_attributes(self):
+        """Return the state attributes."""
+        return self._attributes
+
     @property
     def should_poll(self):
         """No polling needed."""
@@ -229,11 +247,23 @@ async def async_update(self):
             else:
                 self._state = None
                 _LOGGER.error("Could not render template %s: %s", self._name, ex)
-        for property_name, template in (
-            ("_icon", self._icon_template),
-            ("_entity_picture", self._entity_picture_template),
-            ("_name", self._friendly_name_template),
-        ):
+
+        templates = {
+            "_icon": self._icon_template,
+            "_entity_picture": self._entity_picture_template,
+            "_name": self._friendly_name_template,
+        }
+
+        attrs = {}
+        for key, value in self._attribute_templates.items():
+            try:
+                attrs[key] = value.async_render()
+            except TemplateError as err:
+                _LOGGER.error("Error rendering attribute %s: %s", key, err)
+
+        self._attributes = attrs
+
+        for property_name, template in templates.items():
             if template is None:
                 continue
 

tests/components/template/test_sensor.py

@@ -174,6 +174,38 @@ def test_friendly_name_template_with_unknown_state(self):
         state = self.hass.states.get("sensor.test_template_sensor")
         assert state.attributes["friendly_name"] == "It Works."
 
+    def test_attribute_templates(self):
+        """Test attribute_templates template."""
+        with assert_setup_component(1):
+            assert setup_component(
+                self.hass,
+                "sensor",
+                {
+                    "sensor": {
+                        "platform": "template",
+                        "sensors": {
+                            "test_template_sensor": {
+                                "value_template": "{{ states.sensor.test_state.state }}",
+                                "attribute_templates": {
+                                    "test_attribute": "It {{ states.sensor.test_state.state }}."
+                                },
+                            }
+                        },
+                    }
+                },
+            )
+
+        self.hass.start()
+        self.hass.block_till_done()
+
+        state = self.hass.states.get("sensor.test_template_sensor")
+        assert state.attributes.get("test_attribute") == "It ."
+
+        self.hass.states.set("sensor.test_state", "Works")
+        self.hass.block_till_done()
+        state = self.hass.states.get("sensor.test_template_sensor")
+        assert state.attributes["test_attribute"] == "It Works."
+
     def test_template_syntax_error(self):
         """Test templating syntax error."""
         with assert_setup_component(0):
@@ -345,6 +377,34 @@ def test_setup_valid_device_class(self):
         assert "device_class" not in state.attributes
 
 
+async def test_invalid_attribute_template(hass, caplog):
+    """Test that errors are logged if rendering template fails."""
+    hass.states.async_set("sensor.test_sensor", "startup")
+
+    await async_setup_component(
+        hass,
+        "sensor",
+        {
+            "sensor": {
+                "platform": "template",
+                "sensors": {
+                    "invalid_template": {
+                        "value_template": "{{ states.sensor.test_sensor.state }}",
+                        "attribute_templates": {
+                            "test_attribute": "{{ states.sensor.unknown.attributes.picture }}"
+                        },
+                    }
+                },
+            }
+        },
+    )
+    await hass.async_block_till_done()
+    assert len(hass.states.async_all()) == 2
+    await hass.helpers.entity_component.async_update_entity("sensor.invalid_template")
+
+    assert ("Error rendering attribute test_attribute") in caplog.text
+
+
 async def test_no_template_match_all(hass, caplog):
     """Test that we do not allow sensors that match on all."""
     hass.states.async_set("sensor.test_sensor", "startup")
@@ -369,12 +429,22 @@ async def test_no_template_match_all(hass, caplog):
                         "value_template": "{{ states.sensor.test_sensor.state }}",
                         "friendly_name_template": "{{ 1 + 1 }}",
                     },
+                    "invalid_attribute": {
+                        "value_template": "{{ states.sensor.test_sensor.state }}",
+                        "attribute_templates": {"test_attribute": "{{ 1 + 1 }}"},
+                    },
                 },
             }
         },
     )
+
+    assert hass.states.get("sensor.invalid_state").state == "unknown"
+    assert hass.states.get("sensor.invalid_icon").state == "unknown"
+    assert hass.states.get("sensor.invalid_entity_picture").state == "unknown"
+    assert hass.states.get("sensor.invalid_friendly_name").state == "unknown"
+
     await hass.async_block_till_done()
-    assert len(hass.states.async_all()) == 5
+    assert len(hass.states.async_all()) == 6
     assert (
         "Template sensor invalid_state has no entity ids "
         "configured to track nor were we able to extract the entities to "
@@ -395,11 +465,17 @@ async def test_no_template_match_all(hass, caplog):
         "configured to track nor were we able to extract the entities to "
         "track from the friendly_name template"
     ) in caplog.text
+    assert (
+        "Template sensor invalid_attribute has no entity ids "
+        "configured to track nor were we able to extract the entities to "
+        "track from the test_attribute template"
+    ) in caplog.text
 
     assert hass.states.get("sensor.invalid_state").state == "unknown"
     assert hass.states.get("sensor.invalid_icon").state == "unknown"
     assert hass.states.get("sensor.invalid_entity_picture").state == "unknown"
     assert hass.states.get("sensor.invalid_friendly_name").state == "unknown"
+    assert hass.states.get("sensor.invalid_attribute").state == "unknown"
 
     hass.bus.async_fire(EVENT_HOMEASSISTANT_START)
     await hass.async_block_till_done()
@@ -408,6 +484,7 @@ async def test_no_template_match_all(hass, caplog):
     assert hass.states.get("sensor.invalid_icon").state == "startup"
     assert hass.states.get("sensor.invalid_entity_picture").state == "startup"
     assert hass.states.get("sensor.invalid_friendly_name").state == "startup"
+    assert hass.states.get("sensor.invalid_attribute").state == "startup"
 
     hass.states.async_set("sensor.test_sensor", "hello")
     await hass.async_block_till_done()
@@ -416,6 +493,7 @@ async def test_no_template_match_all(hass, caplog):
     assert hass.states.get("sensor.invalid_icon").state == "startup"
     assert hass.states.get("sensor.invalid_entity_picture").state == "startup"
     assert hass.states.get("sensor.invalid_friendly_name").state == "startup"
+    assert hass.states.get("sensor.invalid_attribute").state == "startup"
 
     await hass.helpers.entity_component.async_update_entity("sensor.invalid_state")
     await hass.helpers.entity_component.async_update_entity("sensor.invalid_icon")
@@ -425,8 +503,10 @@ async def test_no_template_match_all(hass, caplog):
     await hass.helpers.entity_component.async_update_entity(
         "sensor.invalid_friendly_name"
     )
+    await hass.helpers.entity_component.async_update_entity("sensor.invalid_attribute")
 
     assert hass.states.get("sensor.invalid_state").state == "2"
     assert hass.states.get("sensor.invalid_icon").state == "hello"
     assert hass.states.get("sensor.invalid_entity_picture").state == "hello"
     assert hass.states.get("sensor.invalid_friendly_name").state == "hello"
+    assert hass.states.get("sensor.invalid_attribute").state == "hello"