core/homeassistant/components/asuswrt/sensor.py

"""Asuswrt status sensors."""
from datetime import timedelta
import enum
import logging
from typing import Any, Dict, List, Optional

from aioasuswrt.asuswrt import AsusWrt

from homeassistant.const import DATA_GIGABYTES, DATA_RATE_MEGABITS_PER_SECOND
from homeassistant.helpers.update_coordinator import (
    CoordinatorEntity,
    DataUpdateCoordinator,
)

from . import DATA_ASUSWRT

UPLOAD_ICON = "mdi:upload-network"
DOWNLOAD_ICON = "mdi:download-network"

_LOGGER = logging.getLogger(__name__)


@enum.unique
class _SensorTypes(enum.Enum):
    DEVICES = "devices"
    UPLOAD = "upload"
    DOWNLOAD = "download"
    DOWNLOAD_SPEED = "download_speed"
    UPLOAD_SPEED = "upload_speed"

    @property
    def unit_of_measurement(self) -> Optional[str]:
        """Return a string with the unit of the sensortype."""
        if self in (_SensorTypes.UPLOAD, _SensorTypes.DOWNLOAD):
            return DATA_GIGABYTES
        if self in (_SensorTypes.UPLOAD_SPEED, _SensorTypes.DOWNLOAD_SPEED):
            return DATA_RATE_MEGABITS_PER_SECOND
        return None

    @property
    def icon(self) -> Optional[str]:
        """Return the expected icon for the sensortype."""
        if self in (_SensorTypes.UPLOAD, _SensorTypes.UPLOAD_SPEED):
            return UPLOAD_ICON
        if self in (_SensorTypes.DOWNLOAD, _SensorTypes.DOWNLOAD_SPEED):
            return DOWNLOAD_ICON
        return None

    @property
    def sensor_name(self) -> Optional[str]:
        """Return the name of the sensor."""
        if self is _SensorTypes.DEVICES:
            return "Asuswrt Devices Connected"
        if self is _SensorTypes.UPLOAD:
            return "Asuswrt Upload"
        if self is _SensorTypes.DOWNLOAD:
            return "Asuswrt Download"
        if self is _SensorTypes.UPLOAD_SPEED:
            return "Asuswrt Upload Speed"
        if self is _SensorTypes.DOWNLOAD_SPEED:
            return "Asuswrt Download Speed"
        return None

    @property
    def is_speed(self) -> bool:
        """Return True if the type is an upload/download speed."""
        return self in (_SensorTypes.UPLOAD_SPEED, _SensorTypes.DOWNLOAD_SPEED)

    @property
    def is_size(self) -> bool:
        """Return True if the type is the total upload/download size."""
        return self in (_SensorTypes.UPLOAD, _SensorTypes.DOWNLOAD)


async def async_setup_platform(hass, config, async_add_entities, discovery_info=None):
    """Set up the asuswrt sensors."""
    if discovery_info is None:
        return

    api: AsusWrt = hass.data[DATA_ASUSWRT]

    # Let's discover the valid sensor types.
    sensors = [_SensorTypes(x) for x in discovery_info]

    data_handler = AsuswrtDataHandler(sensors, api)
    coordinator = DataUpdateCoordinator(
        hass,
        _LOGGER,
        name="sensor",
        update_method=data_handler.update_data,
        # Polling interval. Will only be polled if there are subscribers.
        update_interval=timedelta(seconds=30),
    )

    await coordinator.async_refresh()
    async_add_entities([AsuswrtSensor(coordinator, x) for x in sensors])


class AsuswrtDataHandler:
    """Class handling the API updates."""

    def __init__(self, sensors: List[_SensorTypes], api: AsusWrt):
        """Initialize the handler class."""
        self._api = api
        self._sensors = sensors
        self._connected = True

    async def update_data(self) -> Dict[_SensorTypes, Any]:
        """Fetch the relevant data from the router."""
        ret_dict: Dict[_SensorTypes, Any] = {}
        try:
            if _SensorTypes.DEVICES in self._sensors:
                # Let's check the nr of devices.
                devices = await self._api.async_get_connected_devices()
                ret_dict[_SensorTypes.DEVICES] = len(devices)

            if any(x.is_speed for x in self._sensors):
                # Let's check the upload and download speed
                speed = await self._api.async_get_current_transfer_rates()
                ret_dict[_SensorTypes.DOWNLOAD_SPEED] = round(speed[0] / 125000, 2)
                ret_dict[_SensorTypes.UPLOAD_SPEED] = round(speed[1] / 125000, 2)

            if any(x.is_size for x in self._sensors):
                rates = await self._api.async_get_bytes_total()
                ret_dict[_SensorTypes.DOWNLOAD] = round(rates[0] / 1000000000, 1)
                ret_dict[_SensorTypes.UPLOAD] = round(rates[1] / 1000000000, 1)

            if not self._connected:
                # Log a successful reconnect
                self._connected = True
                _LOGGER.warning("Successfully reconnected to ASUS router")

        except OSError as err:
            if self._connected:
                # Log the first time connection was lost
                _LOGGER.warning("Lost connection to router error due to: '%s'", err)
                self._connected = False

        return ret_dict


class AsuswrtSensor(CoordinatorEntity):
    """The asuswrt specific sensor class."""

    def __init__(self, coordinator: DataUpdateCoordinator, sensor_type: _SensorTypes):
        """Initialize the sensor class."""
        super().__init__(coordinator)
        self._type = sensor_type

    @property
    def state(self):
        """Return the state of the sensor."""
        return self.coordinator.data.get(self._type)

    @property
    def name(self) -> str:
        """Return the name of the sensor."""
        return self._type.sensor_name

    @property
    def icon(self) -> Optional[str]:
        """Return the icon to use in the frontend."""
        return self._type.icon

    @property
    def unit_of_measurement(self) -> Optional[str]:
        """Return the unit of measurement of this entity, if any."""
        return self._type.unit_of_measurement