From 4637a83c29b65804081d98749ffcda6756b18f48 Mon Sep 17 00:00:00 2001
From: Fabian Affolter <fabian@affolter-engineering.ch>
Date: Thu, 10 Mar 2016 20:51:58 +0100
Subject: [PATCH] Add sensor_class to MQTT binary sensor

---
 .../components/binary_sensor/mqtt.py          | 21 +++++++++++---
 tests/components/binary_sensor/test_mqtt.py   | 28 +++++++++++++++++++
 2 files changed, 45 insertions(+), 4 deletions(-)

diff --git a/homeassistant/components/binary_sensor/mqtt.py b/homeassistant/components/binary_sensor/mqtt.py
index 30ea02c9bb4..88cbceccc45 100644
--- a/homeassistant/components/binary_sensor/mqtt.py
+++ b/homeassistant/components/binary_sensor/mqtt.py
@@ -7,7 +7,8 @@ https://home-assistant.io/components/binary_sensor.mqtt/
 import logging
 
 import homeassistant.components.mqtt as mqtt
-from homeassistant.components.binary_sensor import BinarySensorDevice
+from homeassistant.components.binary_sensor import (BinarySensorDevice,
+                                                    SENSOR_CLASSES)
 from homeassistant.const import CONF_VALUE_TEMPLATE
 from homeassistant.helpers import template
 
@@ -28,10 +29,16 @@ def setup_platform(hass, config, add_devices, discovery_info=None):
         _LOGGER.error('Missing required variable: state_topic')
         return False
 
+    sensor_class = config.get('sensor_class')
+    if sensor_class not in SENSOR_CLASSES:
+        _LOGGER.warning('Unknown sensor class: %s', sensor_class)
+        sensor_class = None
+
     add_devices([MqttBinarySensor(
         hass,
         config.get('name', DEFAULT_NAME),
         config.get('state_topic', None),
+        sensor_class,
         config.get('qos', DEFAULT_QOS),
         config.get('payload_on', DEFAULT_PAYLOAD_ON),
         config.get('payload_off', DEFAULT_PAYLOAD_OFF),
@@ -40,15 +47,16 @@ def setup_platform(hass, config, add_devices, discovery_info=None):
 
 # pylint: disable=too-many-arguments, too-many-instance-attributes
 class MqttBinarySensor(BinarySensorDevice):
-    """Represent a binary sensor that is updated by MQTT."""
+    """Representation a binary sensor that is updated by MQTT."""
 
-    def __init__(self, hass, name, state_topic, qos, payload_on, payload_off,
-                 value_template):
+    def __init__(self, hass, name, state_topic, sensor_class, qos, payload_on,
+                 payload_off, value_template):
         """Initialize the MQTT binary sensor."""
         self._hass = hass
         self._name = name
         self._state = False
         self._state_topic = state_topic
+        self._sensor_class = sensor_class
         self._payload_on = payload_on
         self._payload_off = payload_off
         self._qos = qos
@@ -81,3 +89,8 @@ class MqttBinarySensor(BinarySensorDevice):
     def is_on(self):
         """Return true if the binary sensor is on."""
         return self._state
+
+    @property
+    def sensor_class(self):
+        """Return the class of this sensor."""
+        return self._sensor_class
diff --git a/tests/components/binary_sensor/test_mqtt.py b/tests/components/binary_sensor/test_mqtt.py
index 27ecad2e566..75124be10f7 100644
--- a/tests/components/binary_sensor/test_mqtt.py
+++ b/tests/components/binary_sensor/test_mqtt.py
@@ -44,3 +44,31 @@ class TestSensorMQTT(unittest.TestCase):
         self.hass.pool.block_till_done()
         state = self.hass.states.get('binary_sensor.test')
         self.assertEqual(STATE_OFF, state.state)
+
+    def test_valid_sensor_class(self):
+        """Test the setting of a valid sensor class."""
+        self.assertTrue(binary_sensor.setup(self.hass, {
+            'binary_sensor': {
+                'platform': 'mqtt',
+                'name': 'test',
+                'sensor_class': 'motion',
+                'state_topic': 'test-topic',
+            }
+        }))
+
+        state = self.hass.states.get('binary_sensor.test')
+        self.assertEqual('motion', state.attributes.get('sensor_class'))
+
+    def test_invalid_sensor_class(self):
+        """Test the setting of an invalid sensor class."""
+        self.assertTrue(binary_sensor.setup(self.hass, {
+            'binary_sensor': {
+                'platform': 'mqtt',
+                'name': 'test',
+                'sensor_class': 'abc123',
+                'state_topic': 'test-topic',
+            }
+        }))
+
+        state = self.hass.states.get('binary_sensor.test')
+        self.assertIsNone(state.attributes.get('sensor_class'))