From 01ee0661631f0d5a2acaf4707c3c79f457aeee8c Mon Sep 17 00:00:00 2001
From: epenet <6771947+epenet@users.noreply.github.com>
Date: Fri, 9 Dec 2022 13:32:46 +0100
Subject: [PATCH] Use DataRate unit and device class in freebox (#83612)

* Use DataRate unit and device class in freebox

* UnitOfTemperature
---
 homeassistant/components/freebox/const.py  | 35 ++----------------
 homeassistant/components/freebox/sensor.py | 42 +++++++++++++++++++---
 2 files changed, 40 insertions(+), 37 deletions(-)

diff --git a/homeassistant/components/freebox/const.py b/homeassistant/components/freebox/const.py
index 65e5576f9d7..32cf407f2a9 100644
--- a/homeassistant/components/freebox/const.py
+++ b/homeassistant/components/freebox/const.py
@@ -3,8 +3,7 @@ from __future__ import annotations
 
 import socket
 
-from homeassistant.components.sensor import SensorEntityDescription
-from homeassistant.const import DATA_RATE_KILOBYTES_PER_SECOND, PERCENTAGE, Platform
+from homeassistant.const import Platform
 
 DOMAIN = "freebox"
 SERVICE_REBOOT = "reboot"
@@ -26,38 +25,8 @@ STORAGE_KEY = DOMAIN
 STORAGE_VERSION = 1
 
 
-CONNECTION_SENSORS: tuple[SensorEntityDescription, ...] = (
-    SensorEntityDescription(
-        key="rate_down",
-        name="Freebox download speed",
-        native_unit_of_measurement=DATA_RATE_KILOBYTES_PER_SECOND,
-        icon="mdi:download-network",
-    ),
-    SensorEntityDescription(
-        key="rate_up",
-        name="Freebox upload speed",
-        native_unit_of_measurement=DATA_RATE_KILOBYTES_PER_SECOND,
-        icon="mdi:upload-network",
-    ),
-)
-CONNECTION_SENSORS_KEYS: list[str] = [desc.key for desc in CONNECTION_SENSORS]
+CONNECTION_SENSORS_KEYS = {"rate_down", "rate_up"}
 
-CALL_SENSORS: tuple[SensorEntityDescription, ...] = (
-    SensorEntityDescription(
-        key="missed",
-        name="Freebox missed calls",
-        icon="mdi:phone-missed",
-    ),
-)
-
-DISK_PARTITION_SENSORS: tuple[SensorEntityDescription, ...] = (
-    SensorEntityDescription(
-        key="partition_free_space",
-        name="free space",
-        native_unit_of_measurement=PERCENTAGE,
-        icon="mdi:harddisk",
-    ),
-)
 
 # Icons
 DEVICE_ICONS = {
diff --git a/homeassistant/components/freebox/sensor.py b/homeassistant/components/freebox/sensor.py
index e2c42928e80..4d5ba490faf 100644
--- a/homeassistant/components/freebox/sensor.py
+++ b/homeassistant/components/freebox/sensor.py
@@ -10,18 +10,52 @@ from homeassistant.components.sensor import (
     SensorEntityDescription,
 )
 from homeassistant.config_entries import ConfigEntry
-from homeassistant.const import DATA_RATE_KILOBYTES_PER_SECOND, TEMP_CELSIUS
+from homeassistant.const import PERCENTAGE, UnitOfDataRate, UnitOfTemperature
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.dispatcher import async_dispatcher_connect
 from homeassistant.helpers.entity import DeviceInfo
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 import homeassistant.util.dt as dt_util
 
-from .const import CALL_SENSORS, CONNECTION_SENSORS, DISK_PARTITION_SENSORS, DOMAIN
+from .const import DOMAIN
 from .router import FreeboxRouter
 
 _LOGGER = logging.getLogger(__name__)
 
+CONNECTION_SENSORS: tuple[SensorEntityDescription, ...] = (
+    SensorEntityDescription(
+        key="rate_down",
+        name="Freebox download speed",
+        device_class=SensorDeviceClass.DATA_RATE,
+        native_unit_of_measurement=UnitOfDataRate.KILOBYTES_PER_SECOND,
+        icon="mdi:download-network",
+    ),
+    SensorEntityDescription(
+        key="rate_up",
+        name="Freebox upload speed",
+        device_class=SensorDeviceClass.DATA_RATE,
+        native_unit_of_measurement=UnitOfDataRate.KILOBYTES_PER_SECOND,
+        icon="mdi:upload-network",
+    ),
+)
+
+CALL_SENSORS: tuple[SensorEntityDescription, ...] = (
+    SensorEntityDescription(
+        key="missed",
+        name="Freebox missed calls",
+        icon="mdi:phone-missed",
+    ),
+)
+
+DISK_PARTITION_SENSORS: tuple[SensorEntityDescription, ...] = (
+    SensorEntityDescription(
+        key="partition_free_space",
+        name="free space",
+        native_unit_of_measurement=PERCENTAGE,
+        icon="mdi:harddisk",
+    ),
+)
+
 
 async def async_setup_entry(
     hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
@@ -42,7 +76,7 @@ async def async_setup_entry(
             SensorEntityDescription(
                 key=sensor_name,
                 name=f"Freebox {sensor_name}",
-                native_unit_of_measurement=TEMP_CELSIUS,
+                native_unit_of_measurement=UnitOfTemperature.CELSIUS,
                 device_class=SensorDeviceClass.TEMPERATURE,
             ),
         )
@@ -84,7 +118,7 @@ class FreeboxSensor(SensorEntity):
     def async_update_state(self) -> None:
         """Update the Freebox sensor."""
         state = self._router.sensors[self.entity_description.key]
-        if self.native_unit_of_measurement == DATA_RATE_KILOBYTES_PER_SECOND:
+        if self.native_unit_of_measurement == UnitOfDataRate.KILOBYTES_PER_SECOND:
             self._attr_native_value = round(state / 1000, 2)
         else:
             self._attr_native_value = state