From 775751ece5902aebffa89c06baacc471831523a6 Mon Sep 17 00:00:00 2001
From: Erik Montnemery <erik@montnemery.com>
Date: Fri, 6 Oct 2023 10:27:45 +0200
Subject: [PATCH] Add WS command sensor/numeric_device_classes (#101257)

---
 .../components/sensor/websocket_api.py        | 24 ++++++++++++++++-
 tests/components/sensor/test_websocket_api.py | 27 ++++++++++++++++++-
 2 files changed, 49 insertions(+), 2 deletions(-)

diff --git a/homeassistant/components/sensor/websocket_api.py b/homeassistant/components/sensor/websocket_api.py
index 2457bfcabe3..a98c4b25392 100644
--- a/homeassistant/components/sensor/websocket_api.py
+++ b/homeassistant/components/sensor/websocket_api.py
@@ -8,13 +8,19 @@ import voluptuous as vol
 from homeassistant.components import websocket_api
 from homeassistant.core import HomeAssistant, callback
 
-from .const import DEVICE_CLASS_UNITS, UNIT_CONVERTERS
+from .const import (
+    DEVICE_CLASS_UNITS,
+    NON_NUMERIC_DEVICE_CLASSES,
+    UNIT_CONVERTERS,
+    SensorDeviceClass,
+)
 
 
 @callback
 def async_setup(hass: HomeAssistant) -> None:
     """Set up the sensor websocket API."""
     websocket_api.async_register_command(hass, ws_device_class_units)
+    websocket_api.async_register_command(hass, ws_numeric_device_classes)
 
 
 @callback
@@ -36,3 +42,19 @@ def ws_device_class_units(
             key=lambda s: str.casefold(str(s)),
         )
     connection.send_result(msg["id"], {"units": convertible_units})
+
+
+@callback
+@websocket_api.websocket_command(
+    {
+        vol.Required("type"): "sensor/numeric_device_classes",
+    }
+)
+def ws_numeric_device_classes(
+    hass: HomeAssistant, connection: websocket_api.ActiveConnection, msg: dict[str, Any]
+) -> None:
+    """Return numeric sensor device classes."""
+    numeric_device_classes = set(SensorDeviceClass) - NON_NUMERIC_DEVICE_CLASSES
+    connection.send_result(
+        msg["id"], {"numeric_device_classes": list(numeric_device_classes)}
+    )
diff --git a/tests/components/sensor/test_websocket_api.py b/tests/components/sensor/test_websocket_api.py
index 17b8a2ab5cb..bd0a68598e1 100644
--- a/tests/components/sensor/test_websocket_api.py
+++ b/tests/components/sensor/test_websocket_api.py
@@ -1,5 +1,11 @@
 """Test the sensor websocket API."""
-from homeassistant.components.sensor.const import DOMAIN
+from pytest_unordered import unordered
+
+from homeassistant.components.sensor.const import (
+    DOMAIN,
+    NON_NUMERIC_DEVICE_CLASSES,
+    SensorDeviceClass,
+)
 from homeassistant.core import HomeAssistant
 from homeassistant.setup import async_setup_component
 
@@ -59,3 +65,22 @@ async def test_device_class_units(
     msg = await client.receive_json()
     assert msg["success"]
     assert msg["result"] == {"units": []}
+
+
+async def test_numeric_device_classes(
+    hass: HomeAssistant, hass_ws_client: WebSocketGenerator
+) -> None:
+    """Test we can get numeric device classes."""
+    numeric_device_classes = set(SensorDeviceClass) - NON_NUMERIC_DEVICE_CLASSES
+
+    assert await async_setup_component(hass, DOMAIN, {})
+
+    client = await hass_ws_client(hass)
+
+    # Device class with units which sensor allows customizing & converting
+    await client.send_json_auto_id({"type": "sensor/numeric_device_classes"})
+    msg = await client.receive_json()
+    assert msg["success"]
+    assert msg["result"] == {
+        "numeric_device_classes": unordered(list(numeric_device_classes))
+    }