From 2c7bc6eaf8953b222c4f4b99eaf36bd58349b05e Mon Sep 17 00:00:00 2001
From: Erik Eriksson <molobrakos@users.noreply.github.com>
Date: Fri, 23 Mar 2018 11:30:44 +0100
Subject: [PATCH] Support setting icon when configuring MQTT entity (#13304)

---
 homeassistant/components/sensor/mqtt.py | 12 ++++++++++--
 homeassistant/components/switch/mqtt.py | 13 +++++++++++--
 2 files changed, 21 insertions(+), 4 deletions(-)

diff --git a/homeassistant/components/sensor/mqtt.py b/homeassistant/components/sensor/mqtt.py
index b19f5721e4f..d191b9a22e8 100644
--- a/homeassistant/components/sensor/mqtt.py
+++ b/homeassistant/components/sensor/mqtt.py
@@ -17,7 +17,7 @@ from homeassistant.components.mqtt import (
     CONF_PAYLOAD_NOT_AVAILABLE, CONF_QOS, MqttAvailability)
 from homeassistant.const import (
     CONF_FORCE_UPDATE, CONF_NAME, CONF_VALUE_TEMPLATE, STATE_UNKNOWN,
-    CONF_UNIT_OF_MEASUREMENT)
+    CONF_UNIT_OF_MEASUREMENT, CONF_ICON)
 from homeassistant.helpers.entity import Entity
 import homeassistant.components.mqtt as mqtt
 import homeassistant.helpers.config_validation as cv
@@ -36,6 +36,7 @@ DEPENDENCIES = ['mqtt']
 PLATFORM_SCHEMA = mqtt.MQTT_RO_PLATFORM_SCHEMA.extend({
     vol.Optional(CONF_NAME, default=DEFAULT_NAME): cv.string,
     vol.Optional(CONF_UNIT_OF_MEASUREMENT): cv.string,
+    vol.Optional(CONF_ICON): cv.icon,
     vol.Optional(CONF_JSON_ATTRS, default=[]): cv.ensure_list_csv,
     vol.Optional(CONF_EXPIRE_AFTER): cv.positive_int,
     vol.Optional(CONF_FORCE_UPDATE, default=DEFAULT_FORCE_UPDATE): cv.boolean,
@@ -59,6 +60,7 @@ def async_setup_platform(hass, config, async_add_devices, discovery_info=None):
         config.get(CONF_UNIT_OF_MEASUREMENT),
         config.get(CONF_FORCE_UPDATE),
         config.get(CONF_EXPIRE_AFTER),
+        config.get(CONF_ICON),
         value_template,
         config.get(CONF_JSON_ATTRS),
         config.get(CONF_AVAILABILITY_TOPIC),
@@ -71,7 +73,7 @@ class MqttSensor(MqttAvailability, Entity):
     """Representation of a sensor that can be updated using MQTT."""
 
     def __init__(self, name, state_topic, qos, unit_of_measurement,
-                 force_update, expire_after, value_template,
+                 force_update, expire_after, icon, value_template,
                  json_attributes, availability_topic, payload_available,
                  payload_not_available):
         """Initialize the sensor."""
@@ -85,6 +87,7 @@ class MqttSensor(MqttAvailability, Entity):
         self._force_update = force_update
         self._template = value_template
         self._expire_after = expire_after
+        self._icon = icon
         self._expiration_trigger = None
         self._json_attributes = set(json_attributes)
         self._attributes = None
@@ -170,3 +173,8 @@ class MqttSensor(MqttAvailability, Entity):
     def device_state_attributes(self):
         """Return the state attributes."""
         return self._attributes
+
+    @property
+    def icon(self):
+        """Return the icon."""
+        return self._icon
diff --git a/homeassistant/components/switch/mqtt.py b/homeassistant/components/switch/mqtt.py
index a4aea1ded9f..f3bd0bef012 100644
--- a/homeassistant/components/switch/mqtt.py
+++ b/homeassistant/components/switch/mqtt.py
@@ -17,7 +17,7 @@ from homeassistant.components.mqtt import (
 from homeassistant.components.switch import SwitchDevice
 from homeassistant.const import (
     CONF_NAME, CONF_OPTIMISTIC, CONF_VALUE_TEMPLATE, CONF_PAYLOAD_OFF,
-    CONF_PAYLOAD_ON)
+    CONF_PAYLOAD_ON, CONF_ICON)
 import homeassistant.components.mqtt as mqtt
 import homeassistant.helpers.config_validation as cv
 
@@ -32,6 +32,7 @@ DEFAULT_OPTIMISTIC = False
 
 PLATFORM_SCHEMA = mqtt.MQTT_RW_PLATFORM_SCHEMA.extend({
     vol.Optional(CONF_NAME, default=DEFAULT_NAME): cv.string,
+    vol.Optional(CONF_ICON): cv.icon,
     vol.Optional(CONF_PAYLOAD_ON, default=DEFAULT_PAYLOAD_ON): cv.string,
     vol.Optional(CONF_PAYLOAD_OFF, default=DEFAULT_PAYLOAD_OFF): cv.string,
     vol.Optional(CONF_OPTIMISTIC, default=DEFAULT_OPTIMISTIC): cv.boolean,
@@ -50,6 +51,7 @@ def async_setup_platform(hass, config, async_add_devices, discovery_info=None):
 
     async_add_devices([MqttSwitch(
         config.get(CONF_NAME),
+        config.get(CONF_ICON),
         config.get(CONF_STATE_TOPIC),
         config.get(CONF_COMMAND_TOPIC),
         config.get(CONF_AVAILABILITY_TOPIC),
@@ -67,7 +69,8 @@ def async_setup_platform(hass, config, async_add_devices, discovery_info=None):
 class MqttSwitch(MqttAvailability, SwitchDevice):
     """Representation of a switch that can be toggled using MQTT."""
 
-    def __init__(self, name, state_topic, command_topic, availability_topic,
+    def __init__(self, name, icon,
+                 state_topic, command_topic, availability_topic,
                  qos, retain, payload_on, payload_off, optimistic,
                  payload_available, payload_not_available, value_template):
         """Initialize the MQTT switch."""
@@ -75,6 +78,7 @@ class MqttSwitch(MqttAvailability, SwitchDevice):
                          payload_not_available)
         self._state = False
         self._name = name
+        self._icon = icon
         self._state_topic = state_topic
         self._command_topic = command_topic
         self._qos = qos
@@ -130,6 +134,11 @@ class MqttSwitch(MqttAvailability, SwitchDevice):
         """Return true if we do optimistic updates."""
         return self._optimistic
 
+    @property
+    def icon(self):
+        """Return the icon."""
+        return self._icon
+
     @asyncio.coroutine
     def async_turn_on(self, **kwargs):
         """Turn the device on.