From 22af9707ada33eae69c7a93cea09f0e3b83375f0 Mon Sep 17 00:00:00 2001
From: emontnemery <erik@montnemery.com>
Date: Wed, 13 Feb 2019 20:58:46 +0100
Subject: [PATCH] Add support for device_class to MQTT cover (#21044)

---
 homeassistant/components/mqtt/cover.py | 14 +++++++++----
 tests/components/mqtt/test_cover.py    | 28 ++++++++++++++++++++++++++
 2 files changed, 38 insertions(+), 4 deletions(-)

diff --git a/homeassistant/components/mqtt/cover.py b/homeassistant/components/mqtt/cover.py
index 75fae0e9c15..829be266b09 100644
--- a/homeassistant/components/mqtt/cover.py
+++ b/homeassistant/components/mqtt/cover.py
@@ -10,8 +10,8 @@ import voluptuous as vol
 
 from homeassistant.components import cover, mqtt
 from homeassistant.components.cover import (
-    ATTR_POSITION, ATTR_TILT_POSITION, SUPPORT_CLOSE, SUPPORT_CLOSE_TILT,
-    SUPPORT_OPEN, SUPPORT_OPEN_TILT, SUPPORT_SET_POSITION,
+    ATTR_POSITION, ATTR_TILT_POSITION, DEVICE_CLASSES_SCHEMA, SUPPORT_CLOSE,
+    SUPPORT_CLOSE_TILT, SUPPORT_OPEN, SUPPORT_OPEN_TILT, SUPPORT_SET_POSITION,
     SUPPORT_SET_TILT_POSITION, SUPPORT_STOP, SUPPORT_STOP_TILT, CoverDevice)
 from homeassistant.components.mqtt import (
     ATTR_DISCOVERY_HASH, CONF_COMMAND_TOPIC, CONF_QOS, CONF_RETAIN,
@@ -20,8 +20,8 @@ from homeassistant.components.mqtt import (
 from homeassistant.components.mqtt.discovery import (
     MQTT_DISCOVERY_NEW, clear_discovery_hash)
 from homeassistant.const import (
-    CONF_DEVICE, CONF_NAME, CONF_OPTIMISTIC, CONF_VALUE_TEMPLATE, STATE_CLOSED,
-    STATE_OPEN, STATE_UNKNOWN)
+    CONF_DEVICE, CONF_DEVICE_CLASS, CONF_NAME, CONF_OPTIMISTIC,
+    CONF_VALUE_TEMPLATE, STATE_CLOSED, STATE_OPEN, STATE_UNKNOWN)
 from homeassistant.core import callback
 from homeassistant.exceptions import TemplateError
 import homeassistant.helpers.config_validation as cv
@@ -120,6 +120,7 @@ PLATFORM_SCHEMA = vol.All(mqtt.MQTT_BASE_PLATFORM_SCHEMA.extend({
                  default=DEFAULT_TILT_INVERT_STATE): cv.boolean,
     vol.Optional(CONF_UNIQUE_ID): cv.string,
     vol.Optional(CONF_DEVICE): mqtt.MQTT_ENTITY_DEVICE_INFO_SCHEMA,
+    vol.Optional(CONF_DEVICE_CLASS): DEVICE_CLASSES_SCHEMA,
 }).extend(mqtt.MQTT_AVAILABILITY_SCHEMA.schema).extend(
     mqtt.MQTT_JSON_ATTRS_SCHEMA.schema), validate_options)
 
@@ -328,6 +329,11 @@ class MqttCover(MqttAttributes, MqttAvailability, MqttDiscoveryUpdate,
         """Return current position of cover tilt."""
         return self._tilt_value
 
+    @property
+    def device_class(self):
+        """Return the class of this sensor."""
+        return self._config.get(CONF_DEVICE_CLASS)
+
     @property
     def supported_features(self):
         """Flag supported features."""
diff --git a/tests/components/mqtt/test_cover.py b/tests/components/mqtt/test_cover.py
index 343bb3643c6..47681e0de10 100644
--- a/tests/components/mqtt/test_cover.py
+++ b/tests/components/mqtt/test_cover.py
@@ -1051,6 +1051,34 @@ class TestCoverMQTT(unittest.TestCase):
         state = self.hass.states.get('cover.test')
         assert STATE_UNAVAILABLE == state.state
 
+    def test_valid_device_class(self):
+        """Test the setting of a valid sensor class."""
+        assert setup_component(self.hass, cover.DOMAIN, {
+            cover.DOMAIN: {
+                'platform': 'mqtt',
+                'name': 'test',
+                'device_class': 'garage',
+                'state_topic': 'test-topic',
+            }
+        })
+
+        state = self.hass.states.get('cover.test')
+        assert 'garage' == state.attributes.get('device_class')
+
+    def test_invalid_device_class(self):
+        """Test the setting of an invalid sensor class."""
+        assert setup_component(self.hass, cover.DOMAIN, {
+            cover.DOMAIN: {
+                'platform': 'mqtt',
+                'name': 'test',
+                'device_class': 'abc123',
+                'state_topic': 'test-topic',
+            }
+        })
+
+        state = self.hass.states.get('cover.test')
+        assert state is None
+
 
 async def test_setting_attribute_via_mqtt_json_message(hass, mqtt_mock):
     """Test the setting of attribute via MQTT with JSON payload."""