Add 'icon_template' to switch templates (similar to sensor template) (#7862)

* Add 'icon_template' to switch templates (similar to sensor template)

* Add test for template switch 'icon_template'

homeassistant/components/switch/template.py

@@ -25,11 +25,14 @@ from homeassistant.helpers.script import Script
 _LOGGER = logging.getLogger(__name__)
 _VALID_STATES = [STATE_ON, STATE_OFF, 'true', 'false']
 
+CONF_ICON_TEMPLATE = 'icon_template'
+
 ON_ACTION = 'turn_on'
 OFF_ACTION = 'turn_off'
 
 SWITCH_SCHEMA = vol.Schema({
     vol.Required(CONF_VALUE_TEMPLATE): cv.template,
+    vol.Optional(CONF_ICON_TEMPLATE): cv.template,
     vol.Required(ON_ACTION): cv.SCRIPT_SCHEMA,
     vol.Required(OFF_ACTION): cv.SCRIPT_SCHEMA,
     vol.Optional(ATTR_FRIENDLY_NAME): cv.string,
@@ -50,6 +53,7 @@ def async_setup_platform(hass, config, async_add_devices, discovery_info=None):
     for device, device_config in config[CONF_SWITCHES].items():
         friendly_name = device_config.get(ATTR_FRIENDLY_NAME, device)
         state_template = device_config[CONF_VALUE_TEMPLATE]
+        icon_template = device_config.get(CONF_ICON_TEMPLATE)
         on_action = device_config[ON_ACTION]
         off_action = device_config[OFF_ACTION]
         entity_ids = (device_config.get(ATTR_ENTITY_ID) or
@@ -57,10 +61,13 @@ 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
+
         switches.append(
             SwitchTemplate(
-                hass, device, friendly_name, state_template, on_action,
-                off_action, entity_ids)
+                hass, device, friendly_name, state_template, icon_template,
+                on_action, off_action, entity_ids)
             )
     if not switches:
         _LOGGER.error("No switches added")
@@ -74,7 +81,7 @@ class SwitchTemplate(SwitchDevice):
     """Representation of a Template switch."""
 
     def __init__(self, hass, device_id, friendly_name, state_template,
-                 on_action, off_action, entity_ids):
+                 icon_template, on_action, off_action, entity_ids):
         """Initialize the Template switch."""
         self.hass = hass
         self.entity_id = async_generate_entity_id(
@@ -84,6 +91,8 @@ class SwitchTemplate(SwitchDevice):
         self._on_script = Script(hass, on_action)
         self._off_script = Script(hass, off_action)
         self._state = False
+        self._icon_template = icon_template
+        self._icon = None
         self._entities = entity_ids
 
     @asyncio.coroutine
@@ -129,6 +138,11 @@ class SwitchTemplate(SwitchDevice):
         """If switch is available."""
         return self._state is not None
 
+    @property
+    def icon(self):
+        """Return the icon to use in the frontend, if any."""
+        return self._icon
+
     def turn_on(self, **kwargs):
         """Fire the on action."""
         self._on_script.run()
@@ -154,3 +168,17 @@ class SwitchTemplate(SwitchDevice):
         except TemplateError as ex:
             _LOGGER.error(ex)
             self._state = None
+
+        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('Could not render icon template %s,'
+                                    ' the state is unknown.', self._name)
+                    return
+                self._icon = super().icon
+                _LOGGER.error('Could not render icon template %s: %s',
+                              self._name, ex)

tests/components/switch/test_template.py

@@ -130,6 +130,45 @@ class TestTemplateSwitch:
         state = self.hass.states.get('switch.test_template_switch')
         assert state.state == STATE_OFF
 
+    def test_icon_template(self):
+        """Test icon template."""
+        with assert_setup_component(1):
+            assert setup.setup_component(self.hass, 'switch', {
+                'switch': {
+                    'platform': 'template',
+                    'switches': {
+                        'test_template_switch': {
+                            'value_template':
+                                "{{ states.switch.test_state.state }}",
+                            'turn_on': {
+                                'service': 'switch.turn_on',
+                                'entity_id': 'switch.test_state'
+                            },
+                            'turn_off': {
+                                'service': 'switch.turn_off',
+                                'entity_id': 'switch.test_state'
+                            },
+                            'icon_template':
+                                "{% if states.switch.test_state.state %}"
+                                "mdi:check"
+                                "{% endif %}"
+                        }
+                    }
+                }
+            })
+
+        self.hass.start()
+        self.hass.block_till_done()
+
+        state = self.hass.states.get('switch.test_template_switch')
+        assert 'icon' not in state.attributes
+
+        state = self.hass.states.set('switch.test_state', STATE_ON)
+        self.hass.block_till_done()
+
+        state = self.hass.states.get('switch.test_template_switch')
+        assert state.attributes['icon'] == 'mdi:check'
+
     def test_template_syntax_error(self):
         """Test templating syntax error."""
         with assert_setup_component(0):