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

* Sensors implementation and tests improvements

* Remove check for unexpected condition
pull/52230/head
ollo69 2021-06-27 21:09:03 +02:00 committed by GitHub
parent e56069558a
commit e6e39a67f4
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
4 changed files with 70 additions and 109 deletions

View File

@ -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"]

View File

@ -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."""

View File

@ -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]:

View File

@ -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)