core/homeassistant/components/unifi/sensor.py

"""Support for bandwidth sensors with UniFi clients."""
from homeassistant.components.sensor import DEVICE_CLASS_TIMESTAMP, DOMAIN
from homeassistant.const import DATA_MEGABYTES
from homeassistant.core import callback
from homeassistant.helpers.dispatcher import async_dispatcher_connect
import homeassistant.util.dt as dt_util

from .const import DOMAIN as UNIFI_DOMAIN
from .unifi_client import UniFiClient

RX_SENSOR = "rx"
TX_SENSOR = "tx"
UPTIME_SENSOR = "uptime"


async def async_setup_entry(hass, config_entry, async_add_entities):
    """Set up sensors for UniFi integration."""
    controller = hass.data[UNIFI_DOMAIN][config_entry.entry_id]
    controller.entities[DOMAIN] = {
        RX_SENSOR: set(),
        TX_SENSOR: set(),
        UPTIME_SENSOR: set(),
    }

    @callback
    def items_added(
        clients: set = controller.api.clients, devices: set = controller.api.devices
    ) -> None:
        """Update the values of the controller."""
        if controller.option_allow_bandwidth_sensors:
            add_bandwith_entities(controller, async_add_entities, clients)

        if controller.option_allow_uptime_sensors:
            add_uptime_entities(controller, async_add_entities, clients)

    for signal in (controller.signal_update, controller.signal_options_update):
        controller.listeners.append(async_dispatcher_connect(hass, signal, items_added))

    items_added()


@callback
def add_bandwith_entities(controller, async_add_entities, clients):
    """Add new sensor entities from the controller."""
    sensors = []

    for mac in clients:
        for sensor_class in (UniFiRxBandwidthSensor, UniFiTxBandwidthSensor):
            if mac in controller.entities[DOMAIN][sensor_class.TYPE]:
                continue

            client = controller.api.clients[mac]
            sensors.append(sensor_class(client, controller))

    if sensors:
        async_add_entities(sensors)


@callback
def add_uptime_entities(controller, async_add_entities, clients):
    """Add new sensor entities from the controller."""
    sensors = []

    for mac in clients:
        if mac in controller.entities[DOMAIN][UniFiUpTimeSensor.TYPE]:
            continue

        client = controller.api.clients[mac]
        sensors.append(UniFiUpTimeSensor(client, controller))

    if sensors:
        async_add_entities(sensors)


class UniFiBandwidthSensor(UniFiClient):
    """UniFi bandwidth sensor base class."""

    DOMAIN = DOMAIN

    @property
    def name(self) -> str:
        """Return the name of the client."""
        return f"{super().name} {self.TYPE.upper()}"

    @property
    def unit_of_measurement(self) -> str:
        """Return the unit of measurement of this entity."""
        return DATA_MEGABYTES

    async def options_updated(self) -> None:
        """Config entry options are updated, remove entity if option is disabled."""
        if not self.controller.option_allow_bandwidth_sensors:
            await self.remove_item({self.client.mac})


class UniFiRxBandwidthSensor(UniFiBandwidthSensor):
    """Receiving bandwidth sensor."""

    TYPE = RX_SENSOR

    @property
    def state(self) -> int:
        """Return the state of the sensor."""
        if self._is_wired:
            return self.client.wired_rx_bytes / 1000000
        return self.client.rx_bytes / 1000000


class UniFiTxBandwidthSensor(UniFiBandwidthSensor):
    """Transmitting bandwidth sensor."""

    TYPE = TX_SENSOR

    @property
    def state(self) -> int:
        """Return the state of the sensor."""
        if self._is_wired:
            return self.client.wired_tx_bytes / 1000000
        return self.client.tx_bytes / 1000000


class UniFiUpTimeSensor(UniFiClient):
    """UniFi uptime sensor."""

    DOMAIN = DOMAIN
    TYPE = UPTIME_SENSOR

    @property
    def device_class(self) -> str:
        """Return device class."""
        return DEVICE_CLASS_TIMESTAMP

    @property
    def name(self) -> str:
        """Return the name of the client."""
        return f"{super().name} {self.TYPE.capitalize()}"

    @property
    def state(self) -> int:
        """Return the uptime of the client."""
        return dt_util.utc_from_timestamp(float(self.client.uptime)).isoformat()

    async def options_updated(self) -> None:
        """Config entry options are updated, remove entity if option is disabled."""
        if not self.controller.option_allow_uptime_sensors:
            await self.remove_item({self.client.mac})