AsusWRT code improvements for sensors and related tests (#51822)

* Sensors implementation and tests improvements

* Remove check for unexpected condition

homeassistant/components/asuswrt/const.py

@@ -21,8 +21,6 @@ PROTOCOL_SSH = "ssh"
 PROTOCOL_TELNET = "telnet"
 
 # Sensors
-SENSOR_CONNECTED_DEVICE = "sensor_connected_device"
+SENSORS_BYTES = ["sensor_rx_bytes", "sensor_tx_bytes"]
-SENSOR_RX_BYTES = "sensor_rx_bytes"
+SENSORS_CONNECTED_DEVICE = ["sensor_connected_device"]
-SENSOR_TX_BYTES = "sensor_tx_bytes"
+SENSORS_RATES = ["sensor_rx_rates", "sensor_tx_rates"]
-SENSOR_RX_RATES = "sensor_rx_rates"
-SENSOR_TX_RATES = "sensor_tx_rates"

homeassistant/components/asuswrt/router.py

@@ -40,11 +40,9 @@ from .const import (
     DEFAULT_TRACK_UNKNOWN,
     DOMAIN,
     PROTOCOL_TELNET,
-    SENSOR_CONNECTED_DEVICE,
+    SENSORS_BYTES,
-    SENSOR_RX_BYTES,
+    SENSORS_CONNECTED_DEVICE,
-    SENSOR_RX_RATES,
+    SENSORS_RATES,
-    SENSOR_TX_BYTES,
-    SENSOR_TX_RATES,
 )
 
 CONF_REQ_RELOAD = [CONF_DNSMASQ, CONF_INTERFACE, CONF_REQUIRE_IP]
@@ -61,6 +59,16 @@ SENSORS_TYPE_RATES = "sensors_rates"
 _LOGGER = logging.getLogger(__name__)
 
 
+def _get_dict(keys: list, values: list) -> dict[str, Any]:
+    """Create a dict from a list of keys and values."""
+    ret_dict: dict[str, Any] = dict.fromkeys(keys)
+
+    for index, key in enumerate(ret_dict):
+        ret_dict[key] = values[index]
+
+    return ret_dict
+
+
 class AsusWrtSensorDataHandler:
     """Data handler for AsusWrt sensor."""
 
@@ -72,33 +80,25 @@ class AsusWrtSensorDataHandler:
 
     async def _get_connected_devices(self):
         """Return number of connected devices."""
-        return {SENSOR_CONNECTED_DEVICE: self._connected_devices}
+        return {SENSORS_CONNECTED_DEVICE[0]: self._connected_devices}
 
     async def _get_bytes(self):
         """Fetch byte information from the router."""
-        ret_dict: dict[str, Any] = {}
         try:
             datas = await self._api.async_get_bytes_total()
-        except OSError as exc:
+        except (OSError, ValueError) as exc:
-            raise UpdateFailed from exc
+            raise UpdateFailed(exc) from exc
 
-        ret_dict[SENSOR_RX_BYTES] = datas[0]
+        return _get_dict(SENSORS_BYTES, datas)
-        ret_dict[SENSOR_TX_BYTES] = datas[1]
-
-        return ret_dict
 
     async def _get_rates(self):
         """Fetch rates information from the router."""
-        ret_dict: dict[str, Any] = {}
         try:
             rates = await self._api.async_get_current_transfer_rates()
-        except OSError as exc:
+        except (OSError, ValueError) as exc:
-            raise UpdateFailed from exc
+            raise UpdateFailed(exc) from exc
 
-        ret_dict[SENSOR_RX_RATES] = rates[0]
+        return _get_dict(SENSORS_RATES, rates)
-        ret_dict[SENSOR_TX_RATES] = rates[1]
-
-        return ret_dict
 
     def update_device_count(self, conn_devices: int):
         """Update connected devices attribute."""
@@ -315,29 +315,20 @@ class AsusWrtRouter:
         self._sensors_data_handler = AsusWrtSensorDataHandler(self.hass, self._api)
         self._sensors_data_handler.update_device_count(self._connected_devices)
 
-        conn_dev_coordinator = await self._sensors_data_handler.get_coordinator(
+        sensors_types = {
-            SENSORS_TYPE_COUNT, False
+            SENSORS_TYPE_COUNT: SENSORS_CONNECTED_DEVICE,
-        )
+            SENSORS_TYPE_BYTES: SENSORS_BYTES,
-        self._sensors_coordinator[SENSORS_TYPE_COUNT] = {
+            SENSORS_TYPE_RATES: SENSORS_RATES,
-            KEY_COORDINATOR: conn_dev_coordinator,
-            KEY_SENSORS: [SENSOR_CONNECTED_DEVICE],
         }
 
-        bytes_coordinator = await self._sensors_data_handler.get_coordinator(
+        for sensor_type, sensor_names in sensors_types.items():
-            SENSORS_TYPE_BYTES
+            coordinator = await self._sensors_data_handler.get_coordinator(
-        )
+                sensor_type, sensor_type != SENSORS_TYPE_COUNT
-        self._sensors_coordinator[SENSORS_TYPE_BYTES] = {
+            )
-            KEY_COORDINATOR: bytes_coordinator,
+            self._sensors_coordinator[sensor_type] = {
-            KEY_SENSORS: [SENSOR_RX_BYTES, SENSOR_TX_BYTES],
+                KEY_COORDINATOR: coordinator,
-        }
+                KEY_SENSORS: sensor_names,
-
+            }
-        rates_coordinator = await self._sensors_data_handler.get_coordinator(
-            SENSORS_TYPE_RATES
-        )
-        self._sensors_coordinator[SENSORS_TYPE_RATES] = {
-            KEY_COORDINATOR: rates_coordinator,
-            KEY_SENSORS: [SENSOR_RX_RATES, SENSOR_TX_RATES],
-        }
 
     async def _update_unpolled_sensors(self) -> None:
         """Request refresh for AsusWrt unpolled sensors."""

homeassistant/components/asuswrt/sensor.py

@@ -18,11 +18,9 @@ from homeassistant.helpers.update_coordinator import (
 from .const import (
     DATA_ASUSWRT,
     DOMAIN,
-    SENSOR_CONNECTED_DEVICE,
+    SENSORS_BYTES,
-    SENSOR_RX_BYTES,
+    SENSORS_CONNECTED_DEVICE,
-    SENSOR_RX_RATES,
+    SENSORS_RATES,
-    SENSOR_TX_BYTES,
-    SENSOR_TX_RATES,
 )
 from .router import KEY_COORDINATOR, KEY_SENSORS, AsusWrtRouter
 
@@ -38,41 +36,36 @@ SENSOR_DEFAULT_ENABLED = "default_enabled"
 UNIT_DEVICES = "Devices"
 
 CONNECTION_SENSORS = {
-    SENSOR_CONNECTED_DEVICE: {
+    SENSORS_CONNECTED_DEVICE[0]: {
         SENSOR_NAME: "Devices Connected",
         SENSOR_UNIT: UNIT_DEVICES,
         SENSOR_FACTOR: 0,
         SENSOR_ICON: "mdi:router-network",
-        SENSOR_DEVICE_CLASS: None,
         SENSOR_DEFAULT_ENABLED: True,
     },
-    SENSOR_RX_RATES: {
+    SENSORS_RATES[0]: {
         SENSOR_NAME: "Download Speed",
         SENSOR_UNIT: DATA_RATE_MEGABITS_PER_SECOND,
         SENSOR_FACTOR: 125000,
         SENSOR_ICON: "mdi:download-network",
-        SENSOR_DEVICE_CLASS: None,
     },
-    SENSOR_TX_RATES: {
+    SENSORS_RATES[1]: {
         SENSOR_NAME: "Upload Speed",
         SENSOR_UNIT: DATA_RATE_MEGABITS_PER_SECOND,
         SENSOR_FACTOR: 125000,
         SENSOR_ICON: "mdi:upload-network",
-        SENSOR_DEVICE_CLASS: None,
     },
-    SENSOR_RX_BYTES: {
+    SENSORS_BYTES[0]: {
         SENSOR_NAME: "Download",
         SENSOR_UNIT: DATA_GIGABYTES,
         SENSOR_FACTOR: 1000000000,
         SENSOR_ICON: "mdi:download",
-        SENSOR_DEVICE_CLASS: None,
     },
-    SENSOR_TX_BYTES: {
+    SENSORS_BYTES[1]: {
         SENSOR_NAME: "Upload",
         SENSOR_UNIT: DATA_GIGABYTES,
         SENSOR_FACTOR: 1000000000,
         SENSOR_ICON: "mdi:upload",
-        SENSOR_DEVICE_CLASS: None,
     },
 }
 
@@ -108,24 +101,21 @@ class AsusWrtSensor(CoordinatorEntity, SensorEntity):
         coordinator: DataUpdateCoordinator,
         router: AsusWrtRouter,
         sensor_type: str,
-        sensor: dict[str, Any],
+        sensor_def: dict[str, Any],
     ) -> None:
         """Initialize a AsusWrt sensor."""
         super().__init__(coordinator)
         self._router = router
         self._sensor_type = sensor_type
-        self._name = f"{DEFAULT_PREFIX} {sensor[SENSOR_NAME]}"
+        self._sensor_def = sensor_def
+        self._name = f"{DEFAULT_PREFIX} {sensor_def[SENSOR_NAME]}"
         self._unique_id = f"{DOMAIN} {self._name}"
-        self._unit = sensor[SENSOR_UNIT]
+        self._factor = sensor_def.get(SENSOR_FACTOR)
-        self._factor = sensor[SENSOR_FACTOR]
-        self._icon = sensor[SENSOR_ICON]
-        self._device_class = sensor[SENSOR_DEVICE_CLASS]
-        self._default_enabled = sensor.get(SENSOR_DEFAULT_ENABLED, False)
 
     @property
     def entity_registry_enabled_default(self) -> bool:
         """Return if the entity should be enabled when first added to the entity registry."""
-        return self._default_enabled
+        return self._sensor_def.get(SENSOR_DEFAULT_ENABLED, False)
 
     @property
     def state(self) -> str:
@@ -150,17 +140,17 @@ class AsusWrtSensor(CoordinatorEntity, SensorEntity):
     @property
     def unit_of_measurement(self) -> str:
         """Return the unit."""
-        return self._unit
+        return self._sensor_def.get(SENSOR_UNIT)
 
     @property
     def icon(self) -> str:
         """Return the icon."""
-        return self._icon
+        return self._sensor_def.get(SENSOR_ICON)
 
     @property
     def device_class(self) -> str:
         """Return the device_class."""
-        return self._device_class
+        return self._sensor_def.get(SENSOR_DEVICE_CLASS)
 
     @property
     def extra_state_attributes(self) -> dict[str, Any]:

tests/components/asuswrt/test_sensor.py

@@ -20,6 +20,7 @@ from homeassistant.const import (
     STATE_NOT_HOME,
 )
 from homeassistant.helpers import entity_registry as er
+from homeassistant.util import slugify
 from homeassistant.util.dt import utcnow
 
 from tests.common import MockConfigEntry, async_fire_time_changed
@@ -39,6 +40,14 @@ CONFIG_DATA = {
 MOCK_BYTES_TOTAL = [60000000000, 50000000000]
 MOCK_CURRENT_TRANSFER_RATES = [20000000, 10000000]
 
+SENSOR_NAMES = [
+    "Devices Connected",
+    "Download Speed",
+    "Download",
+    "Upload Speed",
+    "Upload",
+]
+
 
 @pytest.fixture(name="mock_devices")
 def mock_devices_fixture():
@@ -88,46 +97,19 @@ async def test_sensors(hass, connect, mock_devices):
 
     # init variable
     unique_id = DOMAIN
-    name_prefix = DEFAULT_PREFIX
+    obj_prefix = slugify(DEFAULT_PREFIX)
-    obj_prefix = name_prefix.lower()
     sensor_prefix = f"{sensor.DOMAIN}.{obj_prefix}"
 
     # Pre-enable the status sensor
-    entity_reg.async_get_or_create(
+    for sensor_name in SENSOR_NAMES:
-        sensor.DOMAIN,
+        sensor_id = slugify(sensor_name)
-        DOMAIN,
+        entity_reg.async_get_or_create(
-        f"{unique_id} {name_prefix} Devices Connected",
+            sensor.DOMAIN,
-        suggested_object_id=f"{obj_prefix}_devices_connected",
+            DOMAIN,
-        disabled_by=None,
+            f"{unique_id} {DEFAULT_PREFIX} {sensor_name}",
-    )
+            suggested_object_id=f"{obj_prefix}_{sensor_id}",
-    entity_reg.async_get_or_create(
+            disabled_by=None,
-        sensor.DOMAIN,
+        )
-        DOMAIN,
-        f"{unique_id} {name_prefix} Download Speed",
-        suggested_object_id=f"{obj_prefix}_download_speed",
-        disabled_by=None,
-    )
-    entity_reg.async_get_or_create(
-        sensor.DOMAIN,
-        DOMAIN,
-        f"{unique_id} {name_prefix} Download",
-        suggested_object_id=f"{obj_prefix}_download",
-        disabled_by=None,
-    )
-    entity_reg.async_get_or_create(
-        sensor.DOMAIN,
-        DOMAIN,
-        f"{unique_id} {name_prefix} Upload Speed",
-        suggested_object_id=f"{obj_prefix}_upload_speed",
-        disabled_by=None,
-    )
-    entity_reg.async_get_or_create(
-        sensor.DOMAIN,
-        DOMAIN,
-        f"{unique_id} {name_prefix} Upload",
-        suggested_object_id=f"{obj_prefix}_upload",
-        disabled_by=None,
-    )
 
     config_entry.add_to_hass(hass)