From 063c0e8f441a51f76487a390260fcf25077081a3 Mon Sep 17 00:00:00 2001
From: Trevor <tboyce021@gmail.com>
Date: Tue, 7 Feb 2017 03:51:44 -0600
Subject: [PATCH] Add icon_template to template sensor (#5766)

* Add icon_template to template sensor

* Update test_template.py

* Update test_template.py again

* Update template.py

* Update test_template.py

* Update test_template.py
---
 homeassistant/components/sensor/template.py | 29 ++++++++++++++++++++-
 tests/components/sensor/test_template.py    | 27 +++++++++++++++++++
 2 files changed, 55 insertions(+), 1 deletion(-)

diff --git a/homeassistant/components/sensor/template.py b/homeassistant/components/sensor/template.py
index 44c6d2e08ea..2535bea1539 100644
--- a/homeassistant/components/sensor/template.py
+++ b/homeassistant/components/sensor/template.py
@@ -21,8 +21,11 @@ import homeassistant.helpers.config_validation as cv
 
 _LOGGER = logging.getLogger(__name__)
 
+CONF_ICON_TEMPLATE = 'icon_template'
+
 SENSOR_SCHEMA = vol.Schema({
     vol.Required(CONF_VALUE_TEMPLATE): cv.template,
+    vol.Optional(CONF_ICON_TEMPLATE): cv.template,
     vol.Optional(ATTR_FRIENDLY_NAME): cv.string,
     vol.Optional(ATTR_UNIT_OF_MEASUREMENT): cv.string,
     vol.Optional(ATTR_ENTITY_ID): cv.entity_ids
@@ -41,6 +44,7 @@ def async_setup_platform(hass, config, async_add_devices, discovery_info=None):
 
     for device, device_config in config[CONF_SENSORS].items():
         state_template = device_config[CONF_VALUE_TEMPLATE]
+        icon_template = device_config.get(CONF_ICON_TEMPLATE)
         entity_ids = (device_config.get(ATTR_ENTITY_ID) or
                       state_template.extract_entities())
         friendly_name = device_config.get(ATTR_FRIENDLY_NAME, device)
@@ -48,6 +52,9 @@ def async_setup_platform(hass, config, async_add_devices, discovery_info=None):
 
         state_template.hass = hass
 
+        if icon_template is not None:
+            icon_template.hass = hass
+
         sensors.append(
             SensorTemplate(
                 hass,
@@ -55,6 +62,7 @@ def async_setup_platform(hass, config, async_add_devices, discovery_info=None):
                 friendly_name,
                 unit_of_measurement,
                 state_template,
+                icon_template,
                 entity_ids)
             )
     if not sensors:
@@ -69,7 +77,7 @@ class SensorTemplate(Entity):
     """Representation of a Template Sensor."""
 
     def __init__(self, hass, device_id, friendly_name, unit_of_measurement,
-                 state_template, entity_ids):
+                 state_template, icon_template, entity_ids):
         """Initialize the sensor."""
         self.hass = hass
         self.entity_id = async_generate_entity_id(ENTITY_ID_FORMAT, device_id,
@@ -78,6 +86,8 @@ class SensorTemplate(Entity):
         self._unit_of_measurement = unit_of_measurement
         self._template = state_template
         self._state = None
+        self._icon_template = icon_template
+        self._icon = None
 
         @callback
         def template_sensor_state_listener(entity, old_state, new_state):
@@ -97,6 +107,11 @@ class SensorTemplate(Entity):
         """Return the state of the sensor."""
         return self._state
 
+    @property
+    def icon(self):
+        """Return the icon to use in the frontend, if any."""
+        return self._icon
+
     @property
     def unit_of_measurement(self):
         """Return the unit_of_measurement of the device."""
@@ -120,3 +135,15 @@ class SensorTemplate(Entity):
                 return
             self._state = None
             _LOGGER.error(ex)
+
+        if self._icon_template is not None:
+            try:
+                self._icon = self._icon_template.async_render()
+            except TemplateError as ex:
+                if ex.args and ex.args[0].startswith(
+                        "UndefinedError: 'None' has no attribute"):
+                    # Common during HA startup - so just a warning
+                    _LOGGER.warning(ex)
+                    return
+                self._icon = super().icon
+                _LOGGER.error(ex)
diff --git a/tests/components/sensor/test_template.py b/tests/components/sensor/test_template.py
index 58f0fb84ac7..0f5e863f328 100644
--- a/tests/components/sensor/test_template.py
+++ b/tests/components/sensor/test_template.py
@@ -41,6 +41,33 @@ class TestTemplateSensor:
         state = self.hass.states.get('sensor.test_template_sensor')
         assert state.state == 'It Works.'
 
+    def test_icon_template(self):
+        """Test icon template."""
+        with assert_setup_component(1):
+            assert setup_component(self.hass, 'sensor', {
+                'sensor': {
+                    'platform': 'template',
+                    'sensors': {
+                        'test_template_sensor': {
+                            'value_template': "State",
+                            'icon_template':
+                                "{% if states.sensor.test_state.state == "
+                                "'Works' %}"
+                                "mdi:check"
+                                "{% endif %}"
+                        }
+                    }
+                }
+            })
+
+        state = self.hass.states.get('sensor.test_template_sensor')
+        assert 'icon' not in state.attributes
+
+        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['icon'] == 'mdi:check'
+
     def test_template_syntax_error(self):
         """Test templating syntax error."""
         with assert_setup_component(0):