diff --git a/homeassistant/components/asuswrt/const.py b/homeassistant/components/asuswrt/const.py index a8977a77ea8..b450030ea3a 100644 --- a/homeassistant/components/asuswrt/const.py +++ b/homeassistant/components/asuswrt/const.py @@ -21,8 +21,6 @@ PROTOCOL_SSH = "ssh" PROTOCOL_TELNET = "telnet" # Sensors -SENSOR_CONNECTED_DEVICE = "sensor_connected_device" -SENSOR_RX_BYTES = "sensor_rx_bytes" -SENSOR_TX_BYTES = "sensor_tx_bytes" -SENSOR_RX_RATES = "sensor_rx_rates" -SENSOR_TX_RATES = "sensor_tx_rates" +SENSORS_BYTES = ["sensor_rx_bytes", "sensor_tx_bytes"] +SENSORS_CONNECTED_DEVICE = ["sensor_connected_device"] +SENSORS_RATES = ["sensor_rx_rates", "sensor_tx_rates"] diff --git a/homeassistant/components/asuswrt/router.py b/homeassistant/components/asuswrt/router.py index 4cea9148470..b94de51b2fb 100644 --- a/homeassistant/components/asuswrt/router.py +++ b/homeassistant/components/asuswrt/router.py @@ -40,11 +40,9 @@ from .const import ( DEFAULT_TRACK_UNKNOWN, DOMAIN, PROTOCOL_TELNET, - SENSOR_CONNECTED_DEVICE, - SENSOR_RX_BYTES, - SENSOR_RX_RATES, - SENSOR_TX_BYTES, - SENSOR_TX_RATES, + SENSORS_BYTES, + SENSORS_CONNECTED_DEVICE, + SENSORS_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: - raise UpdateFailed from exc + except (OSError, ValueError) as exc: + raise UpdateFailed(exc) from exc - ret_dict[SENSOR_RX_BYTES] = datas[0] - ret_dict[SENSOR_TX_BYTES] = datas[1] - - return ret_dict + return _get_dict(SENSORS_BYTES, datas) 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: - raise UpdateFailed from exc + except (OSError, ValueError) as exc: + raise UpdateFailed(exc) from exc - ret_dict[SENSOR_RX_RATES] = rates[0] - ret_dict[SENSOR_TX_RATES] = rates[1] - - return ret_dict + return _get_dict(SENSORS_RATES, rates) 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_TYPE_COUNT, False - ) - self._sensors_coordinator[SENSORS_TYPE_COUNT] = { - KEY_COORDINATOR: conn_dev_coordinator, - KEY_SENSORS: [SENSOR_CONNECTED_DEVICE], + sensors_types = { + SENSORS_TYPE_COUNT: SENSORS_CONNECTED_DEVICE, + SENSORS_TYPE_BYTES: SENSORS_BYTES, + SENSORS_TYPE_RATES: SENSORS_RATES, } - bytes_coordinator = await self._sensors_data_handler.get_coordinator( - SENSORS_TYPE_BYTES - ) - self._sensors_coordinator[SENSORS_TYPE_BYTES] = { - KEY_COORDINATOR: bytes_coordinator, - KEY_SENSORS: [SENSOR_RX_BYTES, SENSOR_TX_BYTES], - } - - 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], - } + for sensor_type, sensor_names in sensors_types.items(): + coordinator = await self._sensors_data_handler.get_coordinator( + sensor_type, sensor_type != SENSORS_TYPE_COUNT + ) + self._sensors_coordinator[sensor_type] = { + KEY_COORDINATOR: coordinator, + KEY_SENSORS: sensor_names, + } async def _update_unpolled_sensors(self) -> None: """Request refresh for AsusWrt unpolled sensors.""" diff --git a/homeassistant/components/asuswrt/sensor.py b/homeassistant/components/asuswrt/sensor.py index 6ec077620f6..cfa8748d2ba 100644 --- a/homeassistant/components/asuswrt/sensor.py +++ b/homeassistant/components/asuswrt/sensor.py @@ -18,11 +18,9 @@ from homeassistant.helpers.update_coordinator import ( from .const import ( DATA_ASUSWRT, DOMAIN, - SENSOR_CONNECTED_DEVICE, - SENSOR_RX_BYTES, - SENSOR_RX_RATES, - SENSOR_TX_BYTES, - SENSOR_TX_RATES, + SENSORS_BYTES, + SENSORS_CONNECTED_DEVICE, + SENSORS_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[SENSOR_FACTOR] - self._icon = sensor[SENSOR_ICON] - self._device_class = sensor[SENSOR_DEVICE_CLASS] - self._default_enabled = sensor.get(SENSOR_DEFAULT_ENABLED, False) + self._factor = sensor_def.get(SENSOR_FACTOR) @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]: diff --git a/tests/components/asuswrt/test_sensor.py b/tests/components/asuswrt/test_sensor.py index 87c3fadb978..cbdded64ff9 100644 --- a/tests/components/asuswrt/test_sensor.py +++ b/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 = name_prefix.lower() + obj_prefix = slugify(DEFAULT_PREFIX) sensor_prefix = f"{sensor.DOMAIN}.{obj_prefix}" # Pre-enable the status sensor - entity_reg.async_get_or_create( - sensor.DOMAIN, - DOMAIN, - f"{unique_id} {name_prefix} Devices Connected", - suggested_object_id=f"{obj_prefix}_devices_connected", - disabled_by=None, - ) - entity_reg.async_get_or_create( - 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, - ) + for sensor_name in SENSOR_NAMES: + sensor_id = slugify(sensor_name) + entity_reg.async_get_or_create( + sensor.DOMAIN, + DOMAIN, + f"{unique_id} {DEFAULT_PREFIX} {sensor_name}", + suggested_object_id=f"{obj_prefix}_{sensor_id}", + disabled_by=None, + ) config_entry.add_to_hass(hass)