From 6ff733b2250f1e92dbc6ec992329d75fc9dcf75b Mon Sep 17 00:00:00 2001
From: Maciej Bieniek <bieniu@users.noreply.github.com>
Date: Fri, 7 Feb 2025 16:36:26 +0100
Subject: [PATCH] Set the device class for the Shelly virtual sensor (#137068)

* Map sensor role to the device class

* Add test

* Fix docstring
---
 homeassistant/components/shelly/const.py  |  6 +++++
 homeassistant/components/shelly/sensor.py | 14 ++++++++++-
 tests/components/shelly/test_sensor.py    | 29 +++++++++++++++++++++++
 3 files changed, 48 insertions(+), 1 deletion(-)

diff --git a/homeassistant/components/shelly/const.py b/homeassistant/components/shelly/const.py
index c8fa72606d6..d47f2b0ae80 100644
--- a/homeassistant/components/shelly/const.py
+++ b/homeassistant/components/shelly/const.py
@@ -28,6 +28,7 @@ from aioshelly.const import (
 )
 
 from homeassistant.components.number import NumberMode
+from homeassistant.components.sensor import SensorDeviceClass
 
 DOMAIN: Final = "shelly"
 
@@ -272,3 +273,8 @@ COMPONENT_ID_PATTERN = re.compile(r"[a-z\d]+:\d+")
 
 # value confirmed by Shelly team
 BLU_TRV_TIMEOUT = 60
+
+ROLE_TO_DEVICE_CLASS_MAP = {
+    "current_humidity": SensorDeviceClass.HUMIDITY,
+    "current_temperature": SensorDeviceClass.TEMPERATURE,
+}
diff --git a/homeassistant/components/shelly/sensor.py b/homeassistant/components/shelly/sensor.py
index 6d000556cf3..c492fc1de9e 100644
--- a/homeassistant/components/shelly/sensor.py
+++ b/homeassistant/components/shelly/sensor.py
@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+from collections.abc import Callable
 from dataclasses import dataclass
 from typing import Final, cast
 
@@ -38,7 +39,7 @@ from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.entity_registry import RegistryEntry
 from homeassistant.helpers.typing import StateType
 
-from .const import CONF_SLEEP_PERIOD, SHAIR_MAX_WORK_HOURS
+from .const import CONF_SLEEP_PERIOD, ROLE_TO_DEVICE_CLASS_MAP, SHAIR_MAX_WORK_HOURS
 from .coordinator import ShellyBlockCoordinator, ShellyConfigEntry, ShellyRpcCoordinator
 from .entity import (
     BlockEntityDescription,
@@ -71,6 +72,8 @@ class BlockSensorDescription(BlockEntityDescription, SensorEntityDescription):
 class RpcSensorDescription(RpcEntityDescription, SensorEntityDescription):
     """Class to describe a RPC sensor."""
 
+    device_class_fn: Callable[[dict], SensorDeviceClass | None] | None = None
+
 
 @dataclass(frozen=True, kw_only=True)
 class RestSensorDescription(RestEntityDescription, SensorEntityDescription):
@@ -95,6 +98,12 @@ class RpcSensor(ShellyRpcAttributeEntity, SensorEntity):
         if self.option_map:
             self._attr_options = list(self.option_map.values())
 
+        if description.device_class_fn is not None:
+            if device_class := description.device_class_fn(
+                coordinator.device.config[key]
+            ):
+                self._attr_device_class = device_class
+
     @property
     def native_value(self) -> StateType:
         """Return value of sensor."""
@@ -1266,6 +1275,9 @@ RPC_SENSORS: Final = {
         unit=lambda config: config["meta"]["ui"]["unit"]
         if config["meta"]["ui"]["unit"]
         else None,
+        device_class_fn=lambda config: ROLE_TO_DEVICE_CLASS_MAP.get(config["role"])
+        if "role" in config
+        else None,
     ),
     "enum": RpcSensorDescription(
         key="enum",
diff --git a/tests/components/shelly/test_sensor.py b/tests/components/shelly/test_sensor.py
index 0bbb374012f..ef7771e53ba 100644
--- a/tests/components/shelly/test_sensor.py
+++ b/tests/components/shelly/test_sensor.py
@@ -1426,3 +1426,32 @@ async def test_blu_trv_sensor_entity(
 
         entry = entity_registry.async_get(entity_id)
         assert entry == snapshot(name=f"{entity_id}-entry")
+
+
+async def test_rpc_device_virtual_number_sensor_with_device_class(
+    hass: HomeAssistant,
+    mock_rpc_device: Mock,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """Test a virtual number sensor with device class for RPC device."""
+    config = deepcopy(mock_rpc_device.config)
+    config["number:203"] = {
+        "name": "Current humidity",
+        "min": 0,
+        "max": 100,
+        "meta": {"ui": {"step": 1, "unit": "%", "view": "label"}},
+        "role": "current_humidity",
+    }
+    monkeypatch.setattr(mock_rpc_device, "config", config)
+
+    status = deepcopy(mock_rpc_device.status)
+    status["number:203"] = {"value": 34}
+    monkeypatch.setattr(mock_rpc_device, "status", status)
+
+    await init_integration(hass, 3)
+
+    state = hass.states.get("sensor.test_name_current_humidity")
+    assert state
+    assert state.state == "34"
+    assert state.attributes.get(ATTR_UNIT_OF_MEASUREMENT) == PERCENTAGE
+    assert state.attributes.get(ATTR_DEVICE_CLASS) == SensorDeviceClass.HUMIDITY