From 589c4ba8113f47d8fa764fb56a8f80f4650a70da Mon Sep 17 00:00:00 2001
From: Pete Ba <pete_ba@yahoo.com>
Date: Mon, 9 Oct 2017 22:21:21 +0100
Subject: [PATCH] Include unit_of_measurement as InfluxDb field

---
 homeassistant/components/influxdb.py |  5 +++-
 tests/components/test_influxdb.py    | 42 ++++++++++++++++++++++++++++
 2 files changed, 46 insertions(+), 1 deletion(-)

diff --git a/homeassistant/components/influxdb.py b/homeassistant/components/influxdb.py
index 1c261d5ec3e07..79749c638f9c5 100644
--- a/homeassistant/components/influxdb.py
+++ b/homeassistant/components/influxdb.py
@@ -149,6 +149,7 @@ def influx_event_listener(event):
             _state = state.state
             _state_key = "state"
 
+        include_uom = True
         measurement = component_config.get(state.entity_id).get(
             CONF_OVERRIDE_MEASUREMENT)
         if measurement in (None, ''):
@@ -161,6 +162,8 @@ def influx_event_listener(event):
                         measurement = default_measurement
                     else:
                         measurement = state.entity_id
+                else:
+                    include_uom = False
 
         json_body = [
             {
@@ -179,7 +182,7 @@ def influx_event_listener(event):
         for key, value in state.attributes.items():
             if key in tags_attributes:
                 json_body[0]['tags'][key] = value
-            elif key != 'unit_of_measurement':
+            elif key != 'unit_of_measurement' or include_uom:
                 # 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 7c98dfcd540df..8815b9eee4a8f 100644
--- a/tests/components/test_influxdb.py
+++ b/tests/components/test_influxdb.py
@@ -532,6 +532,48 @@ def test_event_listener_default_measurement(self, mock_client):
                 self.assertFalse(mock_client.return_value.write_points.called)
             mock_client.return_value.write_points.reset_mock()
 
+    def test_event_listener_unit_of_measurement_field(self, mock_client):
+        """Test the event listener for unit of measurement field."""
+        config = {
+            'influxdb': {
+                'host': 'host',
+                'username': 'user',
+                'password': 'pass',
+                'override_measurement': 'state',
+            }
+        }
+        assert setup_component(self.hass, influxdb.DOMAIN, config)
+        self.handler_method = self.hass.bus.listen.call_args_list[0][0][1]
+
+        attrs = {
+            'unit_of_measurement': 'foobars',
+        }
+        state = mock.MagicMock(
+            state='foo', domain='fake', entity_id='fake.entity-id',
+            object_id='entity', attributes=attrs)
+        event = mock.MagicMock(data={'new_state': state}, time_fired=12345)
+        body = [{
+            'measurement': 'state',
+            'tags': {
+                'domain': 'fake',
+                'entity_id': 'entity',
+            },
+            'time': 12345,
+            'fields': {
+                'state': 'foo',
+                'unit_of_measurement_str': 'foobars',
+            },
+        }]
+        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_tags_attributes(self, mock_client):
         """Test the event listener when some attributes should be tags."""
         config = {