From 16f1ef5a443522e81610811acafc3478726605bd Mon Sep 17 00:00:00 2001 From: Chris Date: Mon, 22 Jun 2020 17:21:20 -0700 Subject: [PATCH] Add attribute templates to template vacuum (#36518) --- homeassistant/components/template/vacuum.py | 30 ++++++++- tests/components/template/test_vacuum.py | 67 +++++++++++++++++++++ 2 files changed, 95 insertions(+), 2 deletions(-) diff --git a/homeassistant/components/template/vacuum.py b/homeassistant/components/template/vacuum.py index c345663ca98..0a1a0e50ffc 100644 --- a/homeassistant/components/template/vacuum.py +++ b/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) - entity_ids = extract_entities(device, "vacuum", None, templates) + initialise_templates(hass, templates, attribute_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 diff --git a/tests/components/template/test_vacuum.py b/tests/components/template/test_vacuum.py index c8ae5bdce51..19cd9f0a8ee 100644 --- a/tests/components/template/test_vacuum.py +++ b/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(