From 944af9cd7d71b00ea1ef519f690e0373da720d39 Mon Sep 17 00:00:00 2001
From: Haim Gelfenbeyn <haim@hageltech.com>
Date: Thu, 3 Aug 2017 10:26:01 -0400
Subject: [PATCH] InfluxDB component improvements (#8633)

* Allow reporting some state attributes as tags to InfluxDB

Some state attributes should really be tags in InfluxDB. E.g.
it is helpful to be able to group by friendly_name, or add a custom
attribute like "location" and group by that. Graphs in Grafana are much
easier to read when friendly names are used, and not node ids.

This commit adds an optional setting to InfluxDB config:
'tags_attributes'. Any attribute on this list will be reported as tag
and not as field to InfluxDB.

* Allow overriding InfluxDB measurement for each reported item separately

Bundling all items with the same "unit of measurement" together does not
always makes sense. For example, both "relatively humidity" and "battery
level" are reported as "%", but I'd rather see them as separate
measurements in InfluxDB. This commit allows for 'influxdb_measurement'
attribute. When set on node, it will take precedence over the global
'override_measurement' and component-specific 'unit_of_measurement'.

* Minor updates to InfluxDB component improvements, as suggested by
@MartinHjelmare.

* Moved per-component config from 'customize' into 'influxdb'
configuration section. The following three sub-sections were added:
'component_config', 'component_config_domain' and
'component_config_glob'. The sole supported per-component attribute
at this point is 'override_measurement'.

* Lint

* Fixed mocked entity_ids in InfluxDB tests to be in domain.entity_id
format, to satisfy EntityValues requirements.

* Added tests for new InfluxDB configuration parameters

* Fixes to some docstrings
---
 homeassistant/components/influxdb.py |  48 +++++++---
 tests/components/test_influxdb.py    | 125 +++++++++++++++++++++++++--
 2 files changed, 153 insertions(+), 20 deletions(-)

diff --git a/homeassistant/components/influxdb.py b/homeassistant/components/influxdb.py
index 6b8cd6734dd..36a58fa8165 100644
--- a/homeassistant/components/influxdb.py
+++ b/homeassistant/components/influxdb.py
@@ -15,6 +15,7 @@ from homeassistant.const import (
     CONF_PORT, CONF_SSL, CONF_VERIFY_SSL, CONF_USERNAME, CONF_PASSWORD,
     CONF_EXCLUDE, CONF_INCLUDE, CONF_DOMAINS, CONF_ENTITIES)
 from homeassistant.helpers import state as state_helper
+from homeassistant.helpers.entity_values import EntityValues
 import homeassistant.helpers.config_validation as cv
 
 REQUIREMENTS = ['influxdb==3.0.0']
@@ -25,13 +26,20 @@ CONF_DB_NAME = 'database'
 CONF_TAGS = 'tags'
 CONF_DEFAULT_MEASUREMENT = 'default_measurement'
 CONF_OVERRIDE_MEASUREMENT = 'override_measurement'
-CONF_BLACKLIST_DOMAINS = "blacklist_domains"
+CONF_TAGS_ATTRIBUTES = 'tags_attributes'
+CONF_COMPONENT_CONFIG = 'component_config'
+CONF_COMPONENT_CONFIG_GLOB = 'component_config_glob'
+CONF_COMPONENT_CONFIG_DOMAIN = 'component_config_domain'
 
 DEFAULT_DATABASE = 'home_assistant'
 DEFAULT_VERIFY_SSL = True
 DOMAIN = 'influxdb'
 TIMEOUT = 5
 
+COMPONENT_CONFIG_SCHEMA_ENTRY = vol.Schema({
+    vol.Optional(CONF_OVERRIDE_MEASUREMENT): cv.string,
+})
+
 CONFIG_SCHEMA = vol.Schema({
     DOMAIN: vol.Schema({
         vol.Optional(CONF_HOST): cv.string,
@@ -54,7 +62,15 @@ CONFIG_SCHEMA = vol.Schema({
         vol.Optional(CONF_OVERRIDE_MEASUREMENT): cv.string,
         vol.Optional(CONF_TAGS, default={}):
             vol.Schema({cv.string: cv.string}),
+        vol.Optional(CONF_TAGS_ATTRIBUTES, default=[]):
+            vol.All(cv.ensure_list, [cv.string]),
         vol.Optional(CONF_VERIFY_SSL, default=DEFAULT_VERIFY_SSL): cv.boolean,
+        vol.Optional(CONF_COMPONENT_CONFIG, default={}):
+            vol.Schema({cv.entity_id: COMPONENT_CONFIG_SCHEMA_ENTRY}),
+        vol.Optional(CONF_COMPONENT_CONFIG_GLOB, default={}):
+            vol.Schema({cv.string: COMPONENT_CONFIG_SCHEMA_ENTRY}),
+        vol.Optional(CONF_COMPONENT_CONFIG_DOMAIN, default={}):
+            vol.Schema({cv.string: COMPONENT_CONFIG_SCHEMA_ENTRY}),
     }),
 }, extra=vol.ALLOW_EXTRA)
 
@@ -96,8 +112,13 @@ def setup(hass, config):
     blacklist_e = set(exclude.get(CONF_ENTITIES, []))
     blacklist_d = set(exclude.get(CONF_DOMAINS, []))
     tags = conf.get(CONF_TAGS)
+    tags_attributes = conf.get(CONF_TAGS_ATTRIBUTES)
     default_measurement = conf.get(CONF_DEFAULT_MEASUREMENT)
     override_measurement = conf.get(CONF_OVERRIDE_MEASUREMENT)
+    component_config = EntityValues(
+        conf[CONF_COMPONENT_CONFIG],
+        conf[CONF_COMPONENT_CONFIG_DOMAIN],
+        conf[CONF_COMPONENT_CONFIG_GLOB])
 
     try:
         influx = InfluxDBClient(**kwargs)
@@ -128,15 +149,18 @@ def setup(hass, config):
             _state = state.state
             _state_key = "state"
 
-        if override_measurement:
-            measurement = override_measurement
-        else:
-            measurement = state.attributes.get('unit_of_measurement')
-            if measurement in (None, ''):
-                if default_measurement:
-                    measurement = default_measurement
-                else:
-                    measurement = state.entity_id
+        measurement = component_config.get(state.entity_id).get(
+            CONF_OVERRIDE_MEASUREMENT)
+        if measurement in (None, ''):
+            if override_measurement:
+                measurement = override_measurement
+            else:
+                measurement = state.attributes.get('unit_of_measurement')
+                if measurement in (None, ''):
+                    if default_measurement:
+                        measurement = default_measurement
+                    else:
+                        measurement = state.entity_id
 
         json_body = [
             {
@@ -153,7 +177,9 @@ def setup(hass, config):
         ]
 
         for key, value in state.attributes.items():
-            if key != 'unit_of_measurement':
+            if key in tags_attributes:
+                json_body[0]['tags'][key] = value
+            elif key != 'unit_of_measurement':
                 # If the key is already in fields
                 if key in json_body[0]['fields']:
                     key = key + "_"
diff --git a/tests/components/test_influxdb.py b/tests/components/test_influxdb.py
index 896ee4699cc..f117b62fddb 100644
--- a/tests/components/test_influxdb.py
+++ b/tests/components/test_influxdb.py
@@ -129,7 +129,8 @@ class TestInfluxDB(unittest.TestCase):
                 'multi_periods': '0.120.240.2023873'
             }
             state = mock.MagicMock(
-                state=in_, domain='fake', object_id='entity', attributes=attrs)
+                state=in_, domain='fake', entity_id='fake.entity-id',
+                object_id='entity', attributes=attrs)
             event = mock.MagicMock(data={'new_state': state}, time_fired=12345)
             if isinstance(out, str):
                 body = [{
@@ -198,11 +199,11 @@ class TestInfluxDB(unittest.TestCase):
             else:
                 attrs = {}
             state = mock.MagicMock(
-                state=1, domain='fake', entity_id='entity-id',
+                state=1, domain='fake', entity_id='fake.entity-id',
                 object_id='entity', attributes=attrs)
             event = mock.MagicMock(data={'new_state': state}, time_fired=12345)
             body = [{
-                'measurement': 'entity-id',
+                'measurement': 'fake.entity-id',
                 'tags': {
                     'domain': 'fake',
                     'entity_id': 'entity',
@@ -227,8 +228,8 @@ class TestInfluxDB(unittest.TestCase):
         self._setup()
 
         state = mock.MagicMock(
-            state=1, domain='fake', entity_id='entity-id', object_id='entity',
-            attributes={})
+            state=1, domain='fake', entity_id='fake.entity-id',
+            object_id='entity', attributes={})
         event = mock.MagicMock(data={'new_state': state}, time_fired=12345)
         mock_client.return_value.write_points.side_effect = \
             influx_client.exceptions.InfluxDBClientError('foo')
@@ -240,11 +241,11 @@ class TestInfluxDB(unittest.TestCase):
 
         for state_state in (1, 'unknown', '', 'unavailable'):
             state = mock.MagicMock(
-                state=state_state, domain='fake', entity_id='entity-id',
+                state=state_state, domain='fake', entity_id='fake.entity-id',
                 object_id='entity', attributes={})
             event = mock.MagicMock(data={'new_state': state}, time_fired=12345)
             body = [{
-                'measurement': 'entity-id',
+                'measurement': 'fake.entity-id',
                 'tags': {
                     'domain': 'fake',
                     'entity_id': 'entity',
@@ -424,7 +425,7 @@ class TestInfluxDB(unittest.TestCase):
             mock_client.return_value.write_points.reset_mock()
 
     def test_event_listener_invalid_type(self, mock_client):
-        """Test the event listener when an attirbute has an invalid type."""
+        """Test the event listener when an attribute has an invalid type."""
         self._setup()
 
         valid = {
@@ -442,7 +443,8 @@ class TestInfluxDB(unittest.TestCase):
                 'invalid_attribute': ['value1', 'value2']
             }
             state = mock.MagicMock(
-                state=in_, domain='fake', object_id='entity', attributes=attrs)
+                state=in_, domain='fake', entity_id='fake.entity-id',
+                object_id='entity', attributes=attrs)
             event = mock.MagicMock(data={'new_state': state}, time_fired=12345)
             if isinstance(out, str):
                 body = [{
@@ -529,3 +531,108 @@ class TestInfluxDB(unittest.TestCase):
             else:
                 self.assertFalse(mock_client.return_value.write_points.called)
             mock_client.return_value.write_points.reset_mock()
+
+    def test_event_listener_tags_attributes(self, mock_client):
+        """Test the event listener when some attributes should be tags."""
+        config = {
+            'influxdb': {
+                'host': 'host',
+                'username': 'user',
+                'password': 'pass',
+                'tags_attributes': ['friendly_fake']
+            }
+        }
+        assert setup_component(self.hass, influxdb.DOMAIN, config)
+        self.handler_method = self.hass.bus.listen.call_args_list[0][0][1]
+
+        attrs = {
+            'friendly_fake': 'tag_str',
+            'field_fake': 'field_str',
+        }
+        state = mock.MagicMock(
+            state=1, domain='fake',
+            entity_id='fake.something',
+            object_id='something', attributes=attrs)
+        event = mock.MagicMock(data={'new_state': state}, time_fired=12345)
+        body = [{
+            'measurement': 'fake.something',
+            'tags': {
+                'domain': 'fake',
+                'entity_id': 'something',
+                'friendly_fake': 'tag_str'
+            },
+            'time': 12345,
+            'fields': {
+                'value': 1,
+                'field_fake_str': 'field_str'
+            },
+        }]
+        self.handler_method(event)
+        self.assertEqual(
+            mock_client.return_value.write_points.call_count, 1
+        )
+        self.assertEqual(
+            mock_client.return_value.write_points.call_args,
+            mock.call(body)
+        )
+        mock_client.return_value.write_points.reset_mock()
+
+    def test_event_listener_component_override_measurement(self, mock_client):
+        """Test the event listener with overrided measurements."""
+        config = {
+            'influxdb': {
+                'host': 'host',
+                'username': 'user',
+                'password': 'pass',
+                'component_config': {
+                    'sensor.fake_humidity': {
+                        'override_measurement': 'humidity'
+                    }
+                },
+                'component_config_glob': {
+                    'binary_sensor.*motion': {
+                        'override_measurement': 'motion'
+                    }
+                },
+                'component_config_domain': {
+                    'climate': {
+                        'override_measurement': 'hvac'
+                    }
+                }
+            }
+        }
+        assert setup_component(self.hass, influxdb.DOMAIN, config)
+        self.handler_method = self.hass.bus.listen.call_args_list[0][0][1]
+
+        test_components = [
+            {'domain': 'sensor', 'id': 'fake_humidity', 'res': 'humidity'},
+            {'domain': 'binary_sensor', 'id': 'fake_motion', 'res': 'motion'},
+            {'domain': 'climate', 'id': 'fake_thermostat', 'res': 'hvac'},
+            {'domain': 'other', 'id': 'just_fake', 'res': 'other.just_fake'},
+        ]
+        for comp in test_components:
+            state = mock.MagicMock(
+                state=1, domain=comp['domain'],
+                entity_id=comp['domain'] + '.' + comp['id'],
+                object_id=comp['id'], attributes={})
+            event = mock.MagicMock(data={'new_state': state}, time_fired=12345)
+            body = [{
+                'measurement': comp['res'],
+                'tags': {
+                    'domain': comp['domain'],
+                    'entity_id': comp['id']
+                },
+                'time': 12345,
+                'fields': {
+                    'value': 1,
+                },
+            }]
+            self.handler_method(event)
+            self.assertEqual(
+                mock_client.return_value.write_points.call_count, 1
+            )
+            self.assertEqual(
+                mock_client.return_value.write_points.call_args,
+                mock.call(body)
+            )
+            mock_client.return_value.write_points.reset_mock()
-- 
GitLab