UniFi - refactor entity management (#34367)
* Move removal of sensor entities into a base class * Fix martins comments on sensors * Reflect sensor changes on device_tracker platform * Reflect sensor changes on switch platform * Improve layering * Make sure to clean up entity and device registry when removing entities * Fix martins comments
This commit is contained in:
Robert Svensson
GitHub
parent
a80ce60e75
commit
e5a861dc90
7 changed files with 252 additions and 381 deletions
|
|
@ -1,6 +1,7 @@
|
|||
"""Support for bandwidth sensors with UniFi clients."""
|
||||
import logging
|
||||
|
||||
from homeassistant.components.sensor import DOMAIN
|
||||
from homeassistant.components.unifi.config_flow import get_controller_from_config_entry
|
||||
from homeassistant.const import DATA_MEGABYTES
|
||||
from homeassistant.core import callback
|
||||
|
|
@ -10,6 +11,9 @@ from .unifi_client import UniFiClient
|
|||
|
||||
LOGGER = logging.getLogger(__name__)
|
||||
|
||||
RX_SENSOR = "rx"
|
||||
TX_SENSOR = "tx"
|
||||
|
||||
|
||||
async def async_setup_platform(hass, config, async_add_entities, discovery_info=None):
|
||||
"""Sensor platform doesn't support configuration through configuration.yaml."""
|
||||
|
|
@ -18,144 +22,74 @@ async def async_setup_platform(hass, config, async_add_entities, discovery_info=
|
|||
async def async_setup_entry(hass, config_entry, async_add_entities):
|
||||
"""Set up sensors for UniFi integration."""
|
||||
controller = get_controller_from_config_entry(hass, config_entry)
|
||||
sensors = {}
|
||||
|
||||
option_allow_bandwidth_sensors = controller.option_allow_bandwidth_sensors
|
||||
|
||||
entity_registry = await hass.helpers.entity_registry.async_get_registry()
|
||||
controller.entities[DOMAIN] = {RX_SENSOR: set(), TX_SENSOR: set()}
|
||||
|
||||
@callback
|
||||
def items_added():
|
||||
"""Update the values of the controller."""
|
||||
nonlocal option_allow_bandwidth_sensors
|
||||
if controller.option_allow_bandwidth_sensors:
|
||||
add_entities(controller, async_add_entities)
|
||||
|
||||
if not option_allow_bandwidth_sensors:
|
||||
return
|
||||
|
||||
add_entities(controller, async_add_entities, sensors)
|
||||
|
||||
controller.listeners.append(
|
||||
async_dispatcher_connect(hass, controller.signal_update, items_added)
|
||||
)
|
||||
|
||||
@callback
|
||||
def items_removed(mac_addresses: set) -> None:
|
||||
"""Items have been removed from the controller."""
|
||||
remove_entities(controller, mac_addresses, sensors, entity_registry)
|
||||
|
||||
controller.listeners.append(
|
||||
async_dispatcher_connect(hass, controller.signal_remove, items_removed)
|
||||
)
|
||||
|
||||
@callback
|
||||
def options_updated():
|
||||
"""Update the values of the controller."""
|
||||
nonlocal option_allow_bandwidth_sensors
|
||||
|
||||
if option_allow_bandwidth_sensors != controller.option_allow_bandwidth_sensors:
|
||||
option_allow_bandwidth_sensors = controller.option_allow_bandwidth_sensors
|
||||
|
||||
if option_allow_bandwidth_sensors:
|
||||
items_added()
|
||||
|
||||
else:
|
||||
for sensor in sensors.values():
|
||||
hass.async_create_task(sensor.async_remove())
|
||||
|
||||
sensors.clear()
|
||||
|
||||
controller.listeners.append(
|
||||
async_dispatcher_connect(
|
||||
hass, controller.signal_options_update, options_updated
|
||||
)
|
||||
)
|
||||
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_entities(controller, async_add_entities, sensors):
|
||||
def add_entities(controller, async_add_entities):
|
||||
"""Add new sensor entities from the controller."""
|
||||
new_sensors = []
|
||||
sensors = []
|
||||
|
||||
for client_id in controller.api.clients:
|
||||
for direction, sensor_class in (
|
||||
("rx", UniFiRxBandwidthSensor),
|
||||
("tx", UniFiTxBandwidthSensor),
|
||||
):
|
||||
item_id = f"{direction}-{client_id}"
|
||||
for mac in controller.api.clients:
|
||||
for sensor_class in (UniFiRxBandwidthSensor, UniFiTxBandwidthSensor):
|
||||
if mac not in controller.entities[DOMAIN][sensor_class.TYPE]:
|
||||
sensors.append(sensor_class(controller.api.clients[mac], controller))
|
||||
|
||||
if item_id in sensors:
|
||||
continue
|
||||
|
||||
sensors[item_id] = sensor_class(
|
||||
controller.api.clients[client_id], controller
|
||||
)
|
||||
new_sensors.append(sensors[item_id])
|
||||
|
||||
if new_sensors:
|
||||
async_add_entities(new_sensors)
|
||||
if sensors:
|
||||
async_add_entities(sensors)
|
||||
|
||||
|
||||
@callback
|
||||
def remove_entities(controller, mac_addresses, sensors, entity_registry):
|
||||
"""Remove select sensor entities."""
|
||||
for mac in mac_addresses:
|
||||
|
||||
for direction in ("rx", "tx"):
|
||||
item_id = f"{direction}-{mac}"
|
||||
|
||||
if item_id not in sensors:
|
||||
continue
|
||||
|
||||
entity = sensors.pop(item_id)
|
||||
controller.hass.async_create_task(entity.async_remove())
|
||||
|
||||
|
||||
class UniFiRxBandwidthSensor(UniFiClient):
|
||||
"""Receiving bandwidth sensor."""
|
||||
class UniFiBandwidthSensor(UniFiClient):
|
||||
"""UniFi bandwidth sensor base class."""
|
||||
|
||||
@property
|
||||
def state(self):
|
||||
"""Return the state of the sensor."""
|
||||
if self._is_wired:
|
||||
return self.client.wired_rx_bytes / 1000000
|
||||
return self.client.raw.get("rx_bytes", 0) / 1000000
|
||||
|
||||
@property
|
||||
def name(self):
|
||||
def name(self) -> str:
|
||||
"""Return the name of the client."""
|
||||
name = self.client.name or self.client.hostname
|
||||
return f"{name} RX"
|
||||
return f"{super().name} {self.TYPE.upper()}"
|
||||
|
||||
@property
|
||||
def unique_id(self):
|
||||
"""Return a unique identifier for this bandwidth sensor."""
|
||||
return f"rx-{self.client.mac}"
|
||||
|
||||
@property
|
||||
def unit_of_measurement(self):
|
||||
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.async_remove()
|
||||
|
||||
class UniFiTxBandwidthSensor(UniFiRxBandwidthSensor):
|
||||
"""Transmitting bandwidth sensor."""
|
||||
|
||||
class UniFiRxBandwidthSensor(UniFiBandwidthSensor):
|
||||
"""Receiving bandwidth sensor."""
|
||||
|
||||
TYPE = RX_SENSOR
|
||||
|
||||
@property
|
||||
def state(self):
|
||||
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.raw.get("tx_bytes", 0) / 1000000
|
||||
|
||||
@property
|
||||
def name(self):
|
||||
"""Return the name of the client."""
|
||||
name = self.client.name or self.client.hostname
|
||||
return f"{name} TX"
|
||||
|
||||
@property
|
||||
def unique_id(self):
|
||||
"""Return a unique identifier for this bandwidth sensor."""
|
||||
return f"tx-{self.client.mac}"
|
||||
return self.client.tx_bytes / 1000000
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue