diff --git a/homeassistant/components/sensor/rest.py b/homeassistant/components/sensor/rest.py index 2ae1c3674ea..86362e8f2d9 100644 --- a/homeassistant/components/sensor/rest.py +++ b/homeassistant/components/sensor/rest.py @@ -5,6 +5,7 @@ For more details about this platform, please refer to the documentation at https://home-assistant.io/components/sensor.rest/ """ import logging +import json import voluptuous as vol import requests @@ -25,6 +26,7 @@ DEFAULT_METHOD = 'GET' DEFAULT_NAME = 'REST Sensor' DEFAULT_VERIFY_SSL = True +CONF_JSON_ATTRS = 'json_attributes' METHODS = ['POST', 'GET'] PLATFORM_SCHEMA = PLATFORM_SCHEMA.extend({ @@ -32,6 +34,7 @@ PLATFORM_SCHEMA = PLATFORM_SCHEMA.extend({ vol.Optional(CONF_AUTHENTICATION): vol.In([HTTP_BASIC_AUTHENTICATION, HTTP_DIGEST_AUTHENTICATION]), vol.Optional(CONF_HEADERS): {cv.string: cv.string}, + vol.Optional(CONF_JSON_ATTRS, default=[]): cv.ensure_list_csv, vol.Optional(CONF_METHOD, default=DEFAULT_METHOD): vol.In(METHODS), vol.Optional(CONF_NAME, default=DEFAULT_NAME): cv.string, vol.Optional(CONF_PASSWORD): cv.string, @@ -55,6 +58,8 @@ def setup_platform(hass, config, add_devices, discovery_info=None): headers = config.get(CONF_HEADERS) unit = config.get(CONF_UNIT_OF_MEASUREMENT) value_template = config.get(CONF_VALUE_TEMPLATE) + json_attrs = config.get(CONF_JSON_ATTRS) + if value_template is not None: value_template.hass = hass @@ -68,13 +73,15 @@ def setup_platform(hass, config, add_devices, discovery_info=None): rest = RestData(method, resource, auth, headers, payload, verify_ssl) rest.update() - add_devices([RestSensor(hass, rest, name, unit, value_template)], True) + add_devices([RestSensor( + hass, rest, name, unit, value_template, json_attrs)], True) class RestSensor(Entity): """Implementation of a REST sensor.""" - def __init__(self, hass, rest, name, unit_of_measurement, value_template): + def __init__(self, hass, rest, name, + unit_of_measurement, value_template, json_attrs): """Initialize the REST sensor.""" self._hass = hass self.rest = rest @@ -82,6 +89,8 @@ class RestSensor(Entity): self._state = STATE_UNKNOWN self._unit_of_measurement = unit_of_measurement self._value_template = value_template + self._json_attrs = json_attrs + self._attributes = None @property def name(self): @@ -108,6 +117,20 @@ class RestSensor(Entity): self.rest.update() value = self.rest.data + if self._json_attrs: + self._attributes = {} + try: + json_dict = json.loads(value) + if isinstance(json_dict, dict): + attrs = {k: json_dict[k] for k in self._json_attrs + if k in json_dict} + self._attributes = attrs + else: + _LOGGER.warning("JSON result was not a dictionary") + except ValueError: + _LOGGER.warning("REST result could not be parsed as JSON") + _LOGGER.debug("Erroneous JSON: %s", value) + if value is None: value = STATE_UNKNOWN elif self._value_template is not None: @@ -116,6 +139,11 @@ class RestSensor(Entity): self._state = value + @property + def device_state_attributes(self): + """Return the state attributes.""" + return self._attributes + class RestData(object): """Class for handling the data retrieval.""" diff --git a/tests/components/sensor/test_rest.py b/tests/components/sensor/test_rest.py index a083dbfb1a2..1bda8ab82f3 100644 --- a/tests/components/sensor/test_rest.py +++ b/tests/components/sensor/test_rest.py @@ -133,9 +133,9 @@ class TestRestSensor(unittest.TestCase): self.value_template = template('{{ value_json.key }}') self.value_template.hass = self.hass - self.sensor = rest.RestSensor( - self.hass, self.rest, self.name, self.unit_of_measurement, - self.value_template) + self.sensor = rest.RestSensor(self.hass, self.rest, self.name, + self.unit_of_measurement, + self.value_template, []) def tearDown(self): """Stop everything that was started.""" @@ -181,12 +181,62 @@ class TestRestSensor(unittest.TestCase): self.rest.update = Mock('rest.RestData.update', side_effect=self.update_side_effect( 'plain_state')) - self.sensor = rest.RestSensor( - self.hass, self.rest, self.name, self.unit_of_measurement, None) + self.sensor = rest.RestSensor(self.hass, self.rest, self.name, + self.unit_of_measurement, None, []) self.sensor.update() self.assertEqual('plain_state', self.sensor.state) self.assertTrue(self.sensor.available) + def test_update_with_json_attrs(self): + """Test attributes get extracted from a JSON result.""" + self.rest.update = Mock('rest.RestData.update', + 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.sensor.update() + self.assertEqual('some_json_value', + self.sensor.device_state_attributes['key']) + + @patch('homeassistant.components.sensor.rest._LOGGER') + def test_update_with_json_attrs_not_dict(self, mock_logger): + """Test attributes get extracted from a JSON result.""" + self.rest.update = Mock('rest.RestData.update', + 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.sensor.update() + self.assertEqual({}, self.sensor.device_state_attributes) + self.assertTrue(mock_logger.warning.called) + + @patch('homeassistant.components.sensor.rest._LOGGER') + def test_update_with_json_attrs_bad_JSON(self, mock_logger): + """Test attributes get extracted from a JSON result.""" + self.rest.update = Mock('rest.RestData.update', + 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.sensor.update() + self.assertEqual({}, self.sensor.device_state_attributes) + self.assertTrue(mock_logger.warning.called) + self.assertTrue(mock_logger.debug.called) + + def test_update_with_json_attrs_and_template(self): + """Test attributes get extracted from a JSON result.""" + self.rest.update = Mock('rest.RestData.update', + side_effect=self.update_side_effect( + '{ "key": "json_state_updated_value" }')) + self.sensor = rest.RestSensor(self.hass, self.rest, self.name, + self.unit_of_measurement, + self.value_template, ['key']) + self.sensor.update() + + self.assertEqual('json_state_updated_value', self.sensor.state) + self.assertEqual('json_state_updated_value', + self.sensor.device_state_attributes['key']) + class TestRestData(unittest.TestCase): """Tests for RestData."""