diff --git a/homeassistant/components/sensor/rest.py b/homeassistant/components/sensor/rest.py index f2fbe6cd191..5b92753eb90 100644 --- a/homeassistant/components/sensor/rest.py +++ b/homeassistant/components/sensor/rest.py @@ -11,12 +11,13 @@ import voluptuous as vol import requests from requests.auth import HTTPBasicAuth, HTTPDigestAuth -from homeassistant.components.sensor import PLATFORM_SCHEMA +from homeassistant.components.sensor import ( + PLATFORM_SCHEMA, DEVICE_CLASSES_SCHEMA) from homeassistant.const import ( CONF_AUTHENTICATION, CONF_FORCE_UPDATE, CONF_HEADERS, CONF_NAME, CONF_METHOD, CONF_PASSWORD, CONF_PAYLOAD, CONF_RESOURCE, CONF_UNIT_OF_MEASUREMENT, CONF_USERNAME, - CONF_VALUE_TEMPLATE, CONF_VERIFY_SSL, + CONF_VALUE_TEMPLATE, CONF_VERIFY_SSL, CONF_DEVICE_CLASS, HTTP_BASIC_AUTHENTICATION, HTTP_DIGEST_AUTHENTICATION, STATE_UNKNOWN) from homeassistant.exceptions import PlatformNotReady from homeassistant.helpers.entity import Entity @@ -43,6 +44,7 @@ PLATFORM_SCHEMA = PLATFORM_SCHEMA.extend({ vol.Optional(CONF_PASSWORD): cv.string, vol.Optional(CONF_PAYLOAD): cv.string, vol.Optional(CONF_UNIT_OF_MEASUREMENT): cv.string, + vol.Optional(CONF_DEVICE_CLASS): DEVICE_CLASSES_SCHEMA, vol.Optional(CONF_USERNAME): cv.string, vol.Optional(CONF_VALUE_TEMPLATE): cv.template, vol.Optional(CONF_VERIFY_SSL, default=DEFAULT_VERIFY_SSL): cv.boolean, @@ -61,6 +63,7 @@ def setup_platform(hass, config, add_entities, discovery_info=None): password = config.get(CONF_PASSWORD) headers = config.get(CONF_HEADERS) unit = config.get(CONF_UNIT_OF_MEASUREMENT) + device_class = config.get(CONF_DEVICE_CLASS) value_template = config.get(CONF_VALUE_TEMPLATE) json_attrs = config.get(CONF_JSON_ATTRS) force_update = config.get(CONF_FORCE_UPDATE) @@ -83,7 +86,8 @@ def setup_platform(hass, config, add_entities, discovery_info=None): # Must update the sensor now (including fetching the rest resource) to # ensure it's updating its state. add_entities([RestSensor( - hass, rest, name, unit, value_template, json_attrs, force_update + hass, rest, name, unit, device_class, + value_template, json_attrs, force_update )], True) @@ -91,13 +95,14 @@ class RestSensor(Entity): """Implementation of a REST sensor.""" def __init__(self, hass, rest, name, unit_of_measurement, - value_template, json_attrs, force_update): + device_class, value_template, json_attrs, force_update): """Initialize the REST sensor.""" self._hass = hass self.rest = rest self._name = name self._state = STATE_UNKNOWN self._unit_of_measurement = unit_of_measurement + self._device_class = device_class self._value_template = value_template self._json_attrs = json_attrs self._attributes = None @@ -113,6 +118,11 @@ class RestSensor(Entity): """Return the unit the value is expressed in.""" return self._unit_of_measurement + @property + def device_class(self): + """Return the class of this sensor.""" + return self._device_class + @property def available(self): """Return if the sensor data are available.""" diff --git a/tests/components/sensor/test_rest.py b/tests/components/sensor/test_rest.py index 2ce72fc4fc4..05ab628b1a8 100644 --- a/tests/components/sensor/test_rest.py +++ b/tests/components/sensor/test_rest.py @@ -135,13 +135,14 @@ class TestRestSensor(unittest.TestCase): '{ "key": "' + self.initial_state + '" }')) self.name = 'foo' self.unit_of_measurement = 'MB' + self.device_class = None self.value_template = template('{{ value_json.key }}') self.value_template.hass = self.hass self.force_update = False self.sensor = rest.RestSensor( self.hass, self.rest, self.name, self.unit_of_measurement, - self.value_template, [], self.force_update + self.device_class, self.value_template, [], self.force_update ) def tearDown(self): @@ -192,7 +193,8 @@ class TestRestSensor(unittest.TestCase): side_effect=self.update_side_effect( 'plain_state')) self.sensor = rest.RestSensor(self.hass, self.rest, self.name, - self.unit_of_measurement, None, [], + self.unit_of_measurement, + self.device_class, None, [], self.force_update) self.sensor.update() assert 'plain_state' == self.sensor.state @@ -204,7 +206,8 @@ class TestRestSensor(unittest.TestCase): side_effect=self.update_side_effect( '{ "key": "some_json_value" }')) self.sensor = rest.RestSensor(self.hass, self.rest, self.name, - self.unit_of_measurement, None, ['key'], + self.unit_of_measurement, + self.device_class, None, ['key'], self.force_update) self.sensor.update() assert 'some_json_value' == \ @@ -216,7 +219,8 @@ class TestRestSensor(unittest.TestCase): self.rest.update = Mock('rest.RestData.update', side_effect=self.update_side_effect(None)) self.sensor = rest.RestSensor(self.hass, self.rest, self.name, - self.unit_of_measurement, None, ['key'], + self.unit_of_measurement, + self.device_class, None, ['key'], self.force_update) self.sensor.update() assert {} == self.sensor.device_state_attributes @@ -229,7 +233,8 @@ class TestRestSensor(unittest.TestCase): side_effect=self.update_side_effect( '["list", "of", "things"]')) self.sensor = rest.RestSensor(self.hass, self.rest, self.name, - self.unit_of_measurement, None, ['key'], + self.unit_of_measurement, + self.device_class, None, ['key'], self.force_update) self.sensor.update() assert {} == self.sensor.device_state_attributes @@ -242,7 +247,8 @@ class TestRestSensor(unittest.TestCase): side_effect=self.update_side_effect( 'This is text rather than JSON data.')) self.sensor = rest.RestSensor(self.hass, self.rest, self.name, - self.unit_of_measurement, None, ['key'], + self.unit_of_measurement, + self.device_class, None, ['key'], self.force_update) self.sensor.update() assert {} == self.sensor.device_state_attributes @@ -256,6 +262,7 @@ class TestRestSensor(unittest.TestCase): '{ "key": "json_state_updated_value" }')) self.sensor = rest.RestSensor(self.hass, self.rest, self.name, self.unit_of_measurement, + self.device_class, self.value_template, ['key'], self.force_update) self.sensor.update()