Command Line Sensor - json_attributes

homeassistant/components/sensor/command_line.py

@@ -5,6 +5,7 @@
 https://home-assistant.io/components/sensor.command_line/
 """
 import logging
+import json
 import subprocess
 import shlex
 
@@ -27,6 +28,7 @@
 
 SCAN_INTERVAL = timedelta(seconds=60)
 
+CONF_JSON_ATTRS = 'json_attributes'
 CONF_COMMAND_TIMEOUT = 'command_timeout'
 DEFAULT_TIMEOUT = 15
 
@@ -35,6 +37,7 @@
     vol.Optional(CONF_NAME, default=DEFAULT_NAME): cv.string,
     vol.Optional(CONF_UNIT_OF_MEASUREMENT): cv.string,
     vol.Optional(CONF_VALUE_TEMPLATE): cv.template,
+    vol.Optional(CONF_JSON_ATTRS, default=[]): cv.ensure_list_csv,
     vol.Optional(
         CONF_COMMAND_TIMEOUT, default=DEFAULT_TIMEOUT): cv.positive_int,
 })
@@ -49,22 +52,27 @@ def setup_platform(hass, config, add_devices, discovery_info=None):
     command_timeout = config.get(CONF_COMMAND_TIMEOUT)
     if value_template is not None:
         value_template.hass = hass
+    json_attrs = config.get(CONF_JSON_ATTRS)
     data = CommandSensorData(hass, command, command_timeout)
 
-    add_devices([CommandSensor(hass, data, name, unit, value_template)], True)
+    add_devices([CommandSensor(hass, data, name, unit, value_template,
+                               json_attrs)], True)
 
 
 class CommandSensor(Entity):
     """Representation of a sensor that is using shell commands."""
 
-    def __init__(self, hass, data, name, unit_of_measurement, value_template):
+    def __init__(self, hass, data, name, unit_of_measurement, value_template,
+                 json_attrs):
         """Initialize the sensor."""
         self._hass = hass
         self.data = data
         self._name = name
         self._state = None
         self._unit_of_measurement = unit_of_measurement
         self._value_template = value_template
+        self._json_attrs = json_attrs
+        self._attributes = None
 
     @property
     def name(self):
@@ -81,11 +89,34 @@ def state(self):
         """Return the state of the device."""
         return self._state
 
+    @property
+    def device_state_attributes(self):
+        """Return the state attributes."""
+        return self._attributes
+
     def update(self):
         """Get the latest data and updates the state."""
         self.data.update()
         value = self.data.value
 
+        if self._json_attrs:
+            self._attributes = {}
+            if value:
+                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("Command result could not be parsed as" +
+                                    "JSON")
+                    _LOGGER.debug("Erroneous JSON: %s", value)
+            else:
+                _LOGGER.warning("Empty reply found when expecting JSON data")
+
         if value is None:
             value = STATE_UNKNOWN
         elif self._value_template is not None:

tests/components/sensor/test_command_line.py

@@ -1,5 +1,6 @@
 """The tests for the Command line sensor platform."""
 import unittest
+from unittest.mock import patch
 
 from homeassistant.helpers.template import Template
 from homeassistant.components.sensor import command_line
@@ -17,6 +18,10 @@ def tearDown(self):
         """Stop everything that was started."""
         self.hass.stop()
 
+    def update_side_effect(self, data):
+        """Side effect function for mocking CommandSensorData.update()."""
+        self.commandline.data = data
+
     def test_setup(self):
         """Test sensor setup."""
         config = {'name': 'Test',
@@ -46,7 +51,7 @@ def test_template(self):
 
         entity = command_line.CommandSensor(
             self.hass, data, 'test', 'in',
-            Template('{{ value | multiply(0.1) }}', self.hass))
+            Template('{{ value | multiply(0.1) }}', self.hass), [])
 
         entity.update()
         self.assertEqual(5, float(entity.state))
@@ -68,3 +73,56 @@ def test_bad_command(self):
         data.update()
 
         self.assertEqual(None, data.value)
+
+    def test_update_with_json_attrs(self):
+        """Test attributes get extracted from a JSON result."""
+        data = command_line.CommandSensorData(
+            self.hass,
+            'echo { \\"key\\": \\"some_json_value\\" }', 15
+        )
+
+        self.sensor = command_line.CommandSensor(self.hass, data, 'test',
+                                                 None, None, ['key'])
+        self.sensor.update()
+        self.assertEqual('some_json_value',
+                         self.sensor.device_state_attributes['key'])
+
+    @patch('homeassistant.components.sensor.command_line._LOGGER')
+    def test_update_with_json_attrs_no_data(self, mock_logger):
+        """Test attributes when no JSON result fetched."""
+        data = command_line.CommandSensorData(
+            self.hass,
+            'echo ', 15
+        )
+        self.sensor = command_line.CommandSensor(self.hass, data, 'test',
+                                                 None, None, ['key'])
+        self.sensor.update()
+        self.assertEqual({}, self.sensor.device_state_attributes)
+        self.assertTrue(mock_logger.warning.called)
+
+    @patch('homeassistant.components.sensor.command_line._LOGGER')
+    def test_update_with_json_attrs_not_dict(self, mock_logger):
+        """Test attributes get extracted from a JSON result."""
+        data = command_line.CommandSensorData(
+            self.hass,
+            'echo [1, 2, 3]', 15
+        )
+        self.sensor = command_line.CommandSensor(self.hass, data, 'test',
+                                                 None, None, ['key'])
+        self.sensor.update()
+        self.assertEqual({}, self.sensor.device_state_attributes)
+        self.assertTrue(mock_logger.warning.called)
+
+    @patch('homeassistant.components.sensor.command_line._LOGGER')
+    def test_update_with_json_attrs_bad_JSON(self, mock_logger):
+        """Test attributes get extracted from a JSON result."""
+        data = command_line.CommandSensorData(
+            self.hass,
+            'echo This is text rather than JSON data.', 15
+        )
+        self.sensor = command_line.CommandSensor(self.hass, data, 'test',
+                                                 None, None, ['key'])
+        self.sensor.update()
+        self.assertEqual({}, self.sensor.device_state_attributes)
+        self.assertTrue(mock_logger.warning.called)
+        self.assertTrue(mock_logger.debug.called)