From 1e57bc5415ba596d5e5e1b2b5880c82c5b455ef9 Mon Sep 17 00:00:00 2001
From: Max <max.toensmann@gmail.com>
Date: Mon, 13 Nov 2023 17:03:08 +0100
Subject: [PATCH] Add `number` state to prometheus metrics (#102518)

---
 .../components/prometheus/__init__.py         | 16 +++--
 tests/components/prometheus/test_init.py      | 65 +++++++++++++++++++
 2 files changed, 76 insertions(+), 5 deletions(-)

diff --git a/homeassistant/components/prometheus/__init__.py b/homeassistant/components/prometheus/__init__.py
index c96ed2e4ed3..1ce16caa6e1 100644
--- a/homeassistant/components/prometheus/__init__.py
+++ b/homeassistant/components/prometheus/__init__.py
@@ -353,18 +353,18 @@ class PrometheusMetrics:
         value = self.state_as_number(state)
         metric.labels(**self._labels(state)).set(value)
 
-    def _handle_input_number(self, state):
+    def _numeric_handler(self, state, domain, title):
         if unit := self._unit_string(state.attributes.get(ATTR_UNIT_OF_MEASUREMENT)):
             metric = self._metric(
-                f"input_number_state_{unit}",
+                f"{domain}_state_{unit}",
                 self.prometheus_cli.Gauge,
-                f"State of the input number measured in {unit}",
+                f"State of the {title} measured in {unit}",
             )
         else:
             metric = self._metric(
-                "input_number_state",
+                f"{domain}_state",
                 self.prometheus_cli.Gauge,
-                "State of the input number",
+                f"State of the {title}",
             )
 
         with suppress(ValueError):
@@ -378,6 +378,12 @@ class PrometheusMetrics:
                 )
             metric.labels(**self._labels(state)).set(value)
 
+    def _handle_input_number(self, state):
+        self._numeric_handler(state, "input_number", "input number")
+
+    def _handle_number(self, state):
+        self._numeric_handler(state, "number", "number")
+
     def _handle_device_tracker(self, state):
         metric = self._metric(
             "device_tracker_state",
diff --git a/tests/components/prometheus/test_init.py b/tests/components/prometheus/test_init.py
index f24782b98d4..1e14ab848a0 100644
--- a/tests/components/prometheus/test_init.py
+++ b/tests/components/prometheus/test_init.py
@@ -20,6 +20,7 @@ from homeassistant.components import (
     input_number,
     light,
     lock,
+    number,
     person,
     prometheus,
     sensor,
@@ -292,6 +293,30 @@ async def test_input_number(client, input_number_entities) -> None:
     )
 
 
+@pytest.mark.parametrize("namespace", [""])
+async def test_number(client, number_entities) -> None:
+    """Test prometheus metrics for number."""
+    body = await generate_latest_metrics(client)
+
+    assert (
+        'number_state{domain="number",'
+        'entity="number.threshold",'
+        'friendly_name="Threshold"} 5.2' in body
+    )
+
+    assert (
+        'number_state{domain="number",'
+        'entity="number.brightness",'
+        'friendly_name="None"} 60.0' in body
+    )
+
+    assert (
+        'number_state_celsius{domain="number",'
+        'entity="number.target_temperature",'
+        'friendly_name="Target temperature"} 22.7' in body
+    )
+
+
 @pytest.mark.parametrize("namespace", [""])
 async def test_battery(client, sensor_entities) -> None:
     """Test prometheus metrics for battery."""
@@ -1382,6 +1407,46 @@ async def input_number_fixture(
     return data
 
 
+@pytest.fixture(name="number_entities")
+async def number_fixture(
+    hass: HomeAssistant, entity_registry: er.EntityRegistry
+) -> dict[str, er.RegistryEntry]:
+    """Simulate number entities."""
+    data = {}
+    number_1 = entity_registry.async_get_or_create(
+        domain=number.DOMAIN,
+        platform="test",
+        unique_id="number_1",
+        suggested_object_id="threshold",
+        original_name="Threshold",
+    )
+    set_state_with_entry(hass, number_1, 5.2)
+    data["number_1"] = number_1
+
+    number_2 = entity_registry.async_get_or_create(
+        domain=number.DOMAIN,
+        platform="test",
+        unique_id="number_2",
+        suggested_object_id="brightness",
+    )
+    set_state_with_entry(hass, number_2, 60)
+    data["number_2"] = number_2
+
+    number_3 = entity_registry.async_get_or_create(
+        domain=number.DOMAIN,
+        platform="test",
+        unique_id="number_3",
+        suggested_object_id="target_temperature",
+        original_name="Target temperature",
+        unit_of_measurement=UnitOfTemperature.CELSIUS,
+    )
+    set_state_with_entry(hass, number_3, 22.7)
+    data["number_3"] = number_3
+
+    await hass.async_block_till_done()
+    return data
+
+
 @pytest.fixture(name="input_boolean_entities")
 async def input_boolean_fixture(
     hass: HomeAssistant, entity_registry: er.EntityRegistry