Unpacking RESTful sensor JSON results into attributes.

homeassistant/components/sensor/rest.py

@@ -5,6 +5,7 @@
 https://home-assistant.io/components/sensor.rest/
 """
 import logging
+import json
 
 import voluptuous as vol
 import requests
@@ -25,13 +26,15 @@
 DEFAULT_NAME = 'REST Sensor'
 DEFAULT_VERIFY_SSL = True
 
+CONF_JSON_ATTRS = 'json_attributes'
 METHODS = ['POST', 'GET']
 
 PLATFORM_SCHEMA = PLATFORM_SCHEMA.extend({
     vol.Required(CONF_RESOURCE): cv.url,
     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,20 +73,24 @@ 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
         self._name = name
         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 @@ def update(self):
         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 @@ def update(self):
 
         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."""

tests/components/sensor/test_rest.py

@@ -133,9 +133,9 @@ def setUp(self):
         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 @@ def test_update_with_no_template(self):
         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."""