Command Line Sensor - json_attributes

homeassistant/components/sensor/command_line.py

@@ -4,39 +4,42 @@
 For more details about this platform, please refer to the documentation at
 https://home-assistant.io/components/sensor.command_line/
 """
+import collections
+from datetime import timedelta
+import json
 import logging
-import subprocess
 import shlex
-
-from datetime import timedelta
+import subprocess
 
 import voluptuous as vol
 
-import homeassistant.helpers.config_validation as cv
 from homeassistant.components.sensor import PLATFORM_SCHEMA
-from homeassistant.helpers import template
-from homeassistant.exceptions import TemplateError
 from homeassistant.const import (
-    CONF_NAME, CONF_VALUE_TEMPLATE, CONF_UNIT_OF_MEASUREMENT, CONF_COMMAND,
+    CONF_COMMAND, CONF_NAME, CONF_UNIT_OF_MEASUREMENT, CONF_VALUE_TEMPLATE,
     STATE_UNKNOWN)
+from homeassistant.exceptions import TemplateError
+from homeassistant.helpers import template
+import homeassistant.helpers.config_validation as cv
 from homeassistant.helpers.entity import Entity
 
 _LOGGER = logging.getLogger(__name__)
 
+CONF_COMMAND_TIMEOUT = 'command_timeout'
+CONF_JSON_ATTRIBUTES = 'json_attributes'
+
 DEFAULT_NAME = 'Command Sensor'
+DEFAULT_TIMEOUT = 15
 
 SCAN_INTERVAL = timedelta(seconds=60)
 
-CONF_COMMAND_TIMEOUT = 'command_timeout'
-DEFAULT_TIMEOUT = 15
-
 PLATFORM_SCHEMA = PLATFORM_SCHEMA.extend({
     vol.Required(CONF_COMMAND): cv.string,
+    vol.Optional(CONF_COMMAND_TIMEOUT, default=DEFAULT_TIMEOUT):
+        cv.positive_int,
+    vol.Optional(CONF_JSON_ATTRIBUTES): cv.ensure_list_csv,
     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_COMMAND_TIMEOUT, default=DEFAULT_TIMEOUT): cv.positive_int,
 })
 
 
@@ -49,18 +52,23 @@ 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_attributes = config.get(CONF_JSON_ATTRIBUTES)
     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_attributes)], 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_attributes):
         """Initialize the sensor."""
         self._hass = hass
         self.data = data
+        self._attributes = None
+        self._json_attributes = json_attributes
         self._name = name
         self._state = None
         self._unit_of_measurement = unit_of_measurement
@@ -81,11 +89,33 @@ 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_attributes:
+            self._attributes = {}
+            if value:
+                try:
+                    json_dict = json.loads(value)
+                    if isinstance(json_dict, collections.Mapping):
+                        self._attributes = {k: json_dict[k] for k in
+                                            self._json_attributes
+                                            if k in json_dict}
+                    else:
+                        _LOGGER.warning("JSON result was not a dictionary")
+                except ValueError:
+                    _LOGGER.warning(
+                        "Unable to parse output as 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,105 @@ 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\\", \\"another_key\\":\
+             \\"another_json_value\\", \\"key_three\\": \\"value_three\\" }'),
+            15
+        )
+
+        self.sensor = command_line.CommandSensor(self.hass, data, 'test',
+                                                 None, None, ['key',
+                                                              'another_key',
+                                                              'key_three'])
+        self.sensor.update()
+        self.assertEqual('some_json_value',
+                         self.sensor.device_state_attributes['key'])
+        self.assertEqual('another_json_value',
+                         self.sensor.device_state_attributes['another_key'])
+        self.assertEqual('value_three',
+                         self.sensor.device_state_attributes['key_three'])
+
+    @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)
+
+    def test_update_with_missing_json_attrs(self):
+        """Test attributes get extracted from a JSON result."""
+        data = command_line.CommandSensorData(
+            self.hass,
+            ('echo { \\"key\\": \\"some_json_value\\", \\"another_key\\":\
+             \\"another_json_value\\", \\"key_three\\": \\"value_three\\" }'),
+            15
+        )
+
+        self.sensor = command_line.CommandSensor(self.hass, data, 'test',
+                                                 None, None, ['key',
+                                                              'another_key',
+                                                              'key_three',
+                                                              'special_key'])
+        self.sensor.update()
+        self.assertEqual('some_json_value',
+                         self.sensor.device_state_attributes['key'])
+        self.assertEqual('another_json_value',
+                         self.sensor.device_state_attributes['another_key'])
+        self.assertEqual('value_three',
+                         self.sensor.device_state_attributes['key_three'])
+        self.assertFalse('special_key' in self.sensor.device_state_attributes)
+
+    def test_update_with_unnecessary_json_attrs(self):
+        """Test attributes get extracted from a JSON result."""
+        data = command_line.CommandSensorData(
+            self.hass,
+            ('echo { \\"key\\": \\"some_json_value\\", \\"another_key\\":\
+             \\"another_json_value\\", \\"key_three\\": \\"value_three\\" }'),
+            15
+        )
+
+        self.sensor = command_line.CommandSensor(self.hass, data, 'test',
+                                                 None, None, ['key',
+                                                              'another_key'])
+        self.sensor.update()
+        self.assertEqual('some_json_value',
+                         self.sensor.device_state_attributes['key'])
+        self.assertEqual('another_json_value',
+                         self.sensor.device_state_attributes['another_key'])
+        self.assertFalse('key_three' in self.sensor.device_state_attributes)