From 027e167d95dd9504e268d2e94b4dd3c5b61057bb Mon Sep 17 00:00:00 2001
From: Franck Nijhof <git@frenck.dev>
Date: Mon, 25 Oct 2021 09:44:43 +0200
Subject: [PATCH] Add support for unit of measurement in MQTT number platform
 (#58343)

---
 homeassistant/components/mqtt/number.py | 15 ++++++++++++++-
 tests/components/mqtt/test_number.py    | 10 +++++++++-
 2 files changed, 23 insertions(+), 2 deletions(-)

diff --git a/homeassistant/components/mqtt/number.py b/homeassistant/components/mqtt/number.py
index a54fd55adb8..f2a472709b3 100644
--- a/homeassistant/components/mqtt/number.py
+++ b/homeassistant/components/mqtt/number.py
@@ -1,4 +1,6 @@
 """Configure number in a device through MQTT topic."""
+from __future__ import annotations
+
 import functools
 import logging
 
@@ -11,7 +13,12 @@ from homeassistant.components.number import (
     DEFAULT_STEP,
     NumberEntity,
 )
-from homeassistant.const import CONF_NAME, CONF_OPTIMISTIC, CONF_VALUE_TEMPLATE
+from homeassistant.const import (
+    CONF_NAME,
+    CONF_OPTIMISTIC,
+    CONF_UNIT_OF_MEASUREMENT,
+    CONF_VALUE_TEMPLATE,
+)
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers import config_validation as cv
 from homeassistant.helpers.reload import async_setup_reload_service
@@ -64,6 +71,7 @@ PLATFORM_SCHEMA = vol.All(
                 vol.Coerce(float), vol.Range(min=1e-3)
             ),
             vol.Optional(CONF_VALUE_TEMPLATE): cv.template,
+            vol.Optional(CONF_UNIT_OF_MEASUREMENT): cv.string,
         },
     ).extend(MQTT_ENTITY_COMMON_SCHEMA.schema),
     validate_config,
@@ -195,6 +203,11 @@ class MqttNumber(MqttEntity, NumberEntity, RestoreEntity):
         """Return the increment/decrement step."""
         return self._config[CONF_STEP]
 
+    @property
+    def unit_of_measurement(self) -> str | None:
+        """Return the unit of measurement."""
+        return self._config.get(CONF_UNIT_OF_MEASUREMENT)
+
     @property
     def value(self):
         """Return the current value."""
diff --git a/tests/components/mqtt/test_number.py b/tests/components/mqtt/test_number.py
index b0989c59ca2..e1e961aa84a 100644
--- a/tests/components/mqtt/test_number.py
+++ b/tests/components/mqtt/test_number.py
@@ -18,7 +18,11 @@ from homeassistant.components.number import (
     DOMAIN as NUMBER_DOMAIN,
     SERVICE_SET_VALUE,
 )
-from homeassistant.const import ATTR_ASSUMED_STATE, ATTR_ENTITY_ID
+from homeassistant.const import (
+    ATTR_ASSUMED_STATE,
+    ATTR_ENTITY_ID,
+    ATTR_UNIT_OF_MEASUREMENT,
+)
 import homeassistant.core as ha
 from homeassistant.setup import async_setup_component
 
@@ -66,6 +70,7 @@ async def test_run_number_setup(hass, mqtt_mock):
                 "state_topic": topic,
                 "command_topic": topic,
                 "name": "Test Number",
+                "unit_of_measurement": "my unit",
                 "payload_reset": "reset!",
             }
         },
@@ -78,6 +83,7 @@ async def test_run_number_setup(hass, mqtt_mock):
 
     state = hass.states.get("number.test_number")
     assert state.state == "10"
+    assert state.attributes.get(ATTR_UNIT_OF_MEASUREMENT) == "my unit"
 
     async_fire_mqtt_message(hass, topic, "20.5")
 
@@ -85,6 +91,7 @@ async def test_run_number_setup(hass, mqtt_mock):
 
     state = hass.states.get("number.test_number")
     assert state.state == "20.5"
+    assert state.attributes.get(ATTR_UNIT_OF_MEASUREMENT) == "my unit"
 
     async_fire_mqtt_message(hass, topic, "reset!")
 
@@ -92,6 +99,7 @@ async def test_run_number_setup(hass, mqtt_mock):
 
     state = hass.states.get("number.test_number")
     assert state.state == "unknown"
+    assert state.attributes.get(ATTR_UNIT_OF_MEASUREMENT) == "my unit"
 
 
 async def test_value_template(hass, mqtt_mock):