Add attribute templates to template vacuum (#36518)

homeassistant/components/template/vacuum.py

@@ -53,6 +53,7 @@ CONF_VACUUMS = "vacuums"
 CONF_BATTERY_LEVEL_TEMPLATE = "battery_level_template"
 CONF_FAN_SPEED_LIST = "fan_speeds"
 CONF_FAN_SPEED_TEMPLATE = "fan_speed_template"
+CONF_ATTRIBUTE_TEMPLATES = "attribute_templates"
 
 ENTITY_ID_FORMAT = DOMAIN + ".{}"
 _VALID_STATES = [
@@ -71,6 +72,9 @@ VACUUM_SCHEMA = vol.Schema(
         vol.Optional(CONF_BATTERY_LEVEL_TEMPLATE): cv.template,
         vol.Optional(CONF_FAN_SPEED_TEMPLATE): cv.template,
         vol.Optional(CONF_AVAILABILITY_TEMPLATE): cv.template,
+        vol.Optional(CONF_ATTRIBUTE_TEMPLATES, default={}): vol.Schema(
+            {cv.string: cv.template}
+        ),
         vol.Required(SERVICE_START): cv.SCRIPT_SCHEMA,
         vol.Optional(SERVICE_PAUSE): cv.SCRIPT_SCHEMA,
         vol.Optional(SERVICE_STOP): cv.SCRIPT_SCHEMA,
@@ -99,6 +103,7 @@ async def async_setup_platform(hass, config, async_add_entities, discovery_info=
         battery_level_template = device_config.get(CONF_BATTERY_LEVEL_TEMPLATE)
         fan_speed_template = device_config.get(CONF_FAN_SPEED_TEMPLATE)
         availability_template = device_config.get(CONF_AVAILABILITY_TEMPLATE)
+        attribute_templates = device_config.get(CONF_ATTRIBUTE_TEMPLATES)
 
         start_action = device_config[SERVICE_START]
         pause_action = device_config.get(SERVICE_PAUSE)
@@ -117,8 +122,10 @@ async def async_setup_platform(hass, config, async_add_entities, discovery_info=
             CONF_AVAILABILITY_TEMPLATE: availability_template,
         }
 
-        initialise_templates(hass, templates)
+        initialise_templates(hass, templates, attribute_templates)
-        entity_ids = extract_entities(device, "vacuum", None, templates)
+        entity_ids = extract_entities(
+            device, "vacuum", None, templates, attribute_templates
+        )
 
         vacuums.append(
             TemplateVacuum(
@@ -138,6 +145,7 @@ async def async_setup_platform(hass, config, async_add_entities, discovery_info=
                 set_fan_speed_action,
                 fan_speed_list,
                 entity_ids,
+                attribute_templates,
             )
         )
 
@@ -165,6 +173,7 @@ class TemplateVacuum(StateVacuumEntity):
         set_fan_speed_action,
         fan_speed_list,
         entity_ids,
+        attribute_templates,
     ):
         """Initialize the vacuum."""
         self.hass = hass
@@ -178,6 +187,8 @@ class TemplateVacuum(StateVacuumEntity):
         self._fan_speed_template = fan_speed_template
         self._availability_template = availability_template
         self._supported_features = SUPPORT_START
+        self._attribute_templates = attribute_templates
+        self._attributes = {}
 
         self._start_script = Script(hass, start_action)
 
@@ -265,6 +276,11 @@ class TemplateVacuum(StateVacuumEntity):
         """Return if the device is available."""
         return self._available
 
+    @property
+    def device_state_attributes(self):
+        """Return the state attributes."""
+        return self._attributes
+
     async def async_start(self):
         """Start or resume the cleaning task."""
         await self._start_script.async_run(context=self._context)
@@ -419,3 +435,13 @@ class TemplateVacuum(StateVacuumEntity):
                     self._name,
                     ex,
                 )
+        # Update attribute if attribute template is defined
+        if self._attribute_templates is not None:
+            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

tests/components/template/test_vacuum.py

@@ -282,6 +282,70 @@ async def test_invalid_availability_template_keeps_component_available(hass, cap
     assert ("UndefinedError: 'x' is undefined") in caplog.text
 
 
+async def test_attribute_templates(hass, calls):
+    """Test attribute_templates template."""
+    assert await setup.async_setup_component(
+        hass,
+        "vacuum",
+        {
+            "vacuum": {
+                "platform": "template",
+                "vacuums": {
+                    "test_template_vacuum": {
+                        "value_template": "{{ 'cleaning' }}",
+                        "start": {"service": "script.vacuum_start"},
+                        "attribute_templates": {
+                            "test_attribute": "It {{ states.sensor.test_state.state }}."
+                        },
+                    }
+                },
+            }
+        },
+    )
+
+    await hass.async_block_till_done()
+    await hass.async_start()
+    await hass.async_block_till_done()
+
+    state = hass.states.get("vacuum.test_template_vacuum")
+    assert state.attributes["test_attribute"] == "It ."
+
+    hass.states.async_set("sensor.test_state", "Works")
+    await hass.async_block_till_done()
+    await hass.helpers.entity_component.async_update_entity(
+        "vacuum.test_template_vacuum"
+    )
+    state = hass.states.get("vacuum.test_template_vacuum")
+    assert state.attributes["test_attribute"] == "It Works."
+
+
+async def test_invalid_attribute_template(hass, caplog):
+    """Test that errors are logged if rendering template fails."""
+    assert await setup.async_setup_component(
+        hass,
+        "vacuum",
+        {
+            "vacuum": {
+                "platform": "template",
+                "vacuums": {
+                    "invalid_template": {
+                        "value_template": "{{ states('input_select.state') }}",
+                        "start": {"service": "script.vacuum_start"},
+                        "attribute_templates": {
+                            "test_attribute": "{{ this_function_does_not_exist() }}"
+                        },
+                    }
+                },
+            }
+        },
+    )
+    await hass.async_block_till_done()
+    assert len(hass.states.async_all()) == 1
+    await hass.helpers.entity_component.async_update_entity("vacuum.invalid_template")
+
+    assert ("Error rendering attribute test_attribute") in caplog.text
+
+
 # End of template tests #
 
 
@@ -529,6 +593,9 @@ async def _register_components(hass):
                 },
             },
             "fan_speeds": ["low", "medium", "high"],
+            "attribute_templates": {
+                "test_attribute": "It {{ states.sensor.test_state.state }}."
+            },
         }
 
         assert await setup.async_setup_component(