From 7f20210e934ab6208dc6b1be46728c35117bb7c3 Mon Sep 17 00:00:00 2001
From: Erik Montnemery <erik@montnemery.com>
Date: Tue, 8 Oct 2019 21:52:25 +0200
Subject: [PATCH] Include unit_of_measurement in sensor device trigger
 capabilities (#27265)

* Expose unit_of_measurement in sensor device trigger

* Update test
---
 .../components/sensor/device_trigger.py       | 16 ++++++++--
 .../components/sensor/test_device_trigger.py  | 31 ++++++++++++++++---
 2 files changed, 41 insertions(+), 6 deletions(-)

diff --git a/homeassistant/components/sensor/device_trigger.py b/homeassistant/components/sensor/device_trigger.py
index bd53dca0c9d..7eabc457161 100644
--- a/homeassistant/components/sensor/device_trigger.py
+++ b/homeassistant/components/sensor/device_trigger.py
@@ -11,6 +11,7 @@ from homeassistant.const import (
     CONF_ENTITY_ID,
     CONF_FOR,
     CONF_TYPE,
+    CONF_UNIT_OF_MEASUREMENT,
     DEVICE_CLASS_BATTERY,
     DEVICE_CLASS_HUMIDITY,
     DEVICE_CLASS_ILLUMINANCE,
@@ -149,11 +150,22 @@ async def async_get_triggers(hass, device_id):
 
 async def async_get_trigger_capabilities(hass, config):
     """List trigger capabilities."""
+    state = hass.states.get(config[CONF_ENTITY_ID])
+    unit_of_measurement = (
+        state.attributes.get(ATTR_UNIT_OF_MEASUREMENT) if state else ""
+    )
+
     return {
         "extra_fields": vol.Schema(
             {
-                vol.Optional(CONF_ABOVE): vol.Coerce(float),
-                vol.Optional(CONF_BELOW): vol.Coerce(float),
+                vol.Optional(
+                    CONF_ABOVE,
+                    description={CONF_UNIT_OF_MEASUREMENT: unit_of_measurement},
+                ): vol.Coerce(float),
+                vol.Optional(
+                    CONF_BELOW,
+                    description={CONF_UNIT_OF_MEASUREMENT: unit_of_measurement},
+                ): vol.Coerce(float),
                 vol.Optional(CONF_FOR): cv.positive_time_period_dict,
             }
         )
diff --git a/tests/components/sensor/test_device_trigger.py b/tests/components/sensor/test_device_trigger.py
index 45452dc84a0..1bc7e5e1ee5 100644
--- a/tests/components/sensor/test_device_trigger.py
+++ b/tests/components/sensor/test_device_trigger.py
@@ -74,26 +74,49 @@ async def test_get_triggers(hass, device_reg, entity_reg):
         if device_class != "none"
     ]
     triggers = await async_get_device_automations(hass, "trigger", device_entry.id)
+    assert len(triggers) == 8
     assert triggers == expected_triggers
 
 
 async def test_get_trigger_capabilities(hass, device_reg, entity_reg):
-    """Test we get the expected capabilities from a binary_sensor trigger."""
+    """Test we get the expected capabilities from a sensor trigger."""
+    platform = getattr(hass.components, f"test.{DOMAIN}")
+    platform.init()
+
     config_entry = MockConfigEntry(domain="test", data={})
     config_entry.add_to_hass(hass)
     device_entry = device_reg.async_get_or_create(
         config_entry_id=config_entry.entry_id,
         connections={(device_registry.CONNECTION_NETWORK_MAC, "12:34:56:AB:CD:EF")},
     )
-    entity_reg.async_get_or_create(DOMAIN, "test", "5678", device_id=device_entry.id)
+    entity_reg.async_get_or_create(
+        DOMAIN,
+        "test",
+        platform.ENTITIES["battery"].unique_id,
+        device_id=device_entry.id,
+    )
+
+    assert await async_setup_component(hass, DOMAIN, {DOMAIN: {CONF_PLATFORM: "test"}})
+
     expected_capabilities = {
         "extra_fields": [
-            {"name": "above", "optional": True, "type": "float"},
-            {"name": "below", "optional": True, "type": "float"},
+            {
+                "description": {"unit_of_measurement": "%"},
+                "name": "above",
+                "optional": True,
+                "type": "float",
+            },
+            {
+                "description": {"unit_of_measurement": "%"},
+                "name": "below",
+                "optional": True,
+                "type": "float",
+            },
             {"name": "for", "optional": True, "type": "positive_time_period_dict"},
         ]
     }
     triggers = await async_get_device_automations(hass, "trigger", device_entry.id)
+    assert len(triggers) == 1
     for trigger in triggers:
         capabilities = await async_get_device_automation_capabilities(
             hass, "trigger", trigger