Home-Assistant
/
core
mirror of https://github.com/home-assistant/core.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

UniFi - Bandwidth sensors (#27229) 

* First draft of UniFi bandwidth sensors

* Clean up

* Add tests for sensors
pull/27438/head
Robert Svensson 2019-10-06 17:17:40 +02:00 committed by Paulus Schoutsen
parent 6cc71db385
commit c7c88b2b68
10 changed files with 444 additions and 12 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
homeassistant/components/unifi/.translations/en.json
View File

@ -32,6 +32,11 @@
"track_devices": "Track network devices (Ubiquiti devices)",
"track_wired_clients": "Include wired network clients"
}
},
"statistics_sensors": {
"data": {
"allow_bandwidth_sensors": "Create bandwidth usage sensors for network clients"
}
}
}
}

31
homeassistant/components/unifi/config_flow.py
View File

@ -12,6 +12,7 @@ from homeassistant.const import (
)
from .const import (
CONF_ALLOW_BANDWIDTH_SENSORS,
CONF_CONTROLLER,
CONF_DETECTION_TIME,
CONF_SITE_ID,
@ -19,6 +20,7 @@ from .const import (
CONF_TRACK_DEVICES,
CONF_TRACK_WIRED_CLIENTS,
CONTROLLER_ID,
DEFAULT_ALLOW_BANDWIDTH_SENSORS,
DEFAULT_TRACK_CLIENTS,
DEFAULT_TRACK_DEVICES,
DEFAULT_TRACK_WIRED_CLIENTS,
@ -171,6 +173,7 @@ class UnifiOptionsFlowHandler(config_entries.OptionsFlow):
def __init__(self, config_entry):
"""Initialize UniFi options flow."""
self.config_entry = config_entry
self.options = dict(config_entry.options)
async def async_step_init(self, user_input=None):
"""Manage the UniFi options."""
@ -179,7 +182,8 @@ class UnifiOptionsFlowHandler(config_entries.OptionsFlow):
async def async_step_device_tracker(self, user_input=None):
"""Manage the device tracker options."""
if user_input is not None:
return self.async_create_entry(title="", data=user_input)
self.options.update(user_input)
return await self.async_step_statistics_sensors()
return self.async_show_form(
step_id="device_tracker",
@ -212,3 +216,28 @@ class UnifiOptionsFlowHandler(config_entries.OptionsFlow):
}
),
)
async def async_step_statistics_sensors(self, user_input=None):
"""Manage the statistics sensors options."""
if user_input is not None:
self.options.update(user_input)
return await self._update_options()
return self.async_show_form(
step_id="statistics_sensors",
data_schema=vol.Schema(
{
vol.Optional(
CONF_ALLOW_BANDWIDTH_SENSORS,
default=self.config_entry.options.get(
CONF_ALLOW_BANDWIDTH_SENSORS,
DEFAULT_ALLOW_BANDWIDTH_SENSORS,
),
): bool
}
),
)
async def _update_options(self):
"""Update config entry options."""
return self.async_create_entry(title="", data=self.options)

2
homeassistant/components/unifi/const.py
View File

@ -12,6 +12,7 @@ CONF_SITE_ID = "site"
UNIFI_CONFIG = "unifi_config"
UNIFI_WIRELESS_CLIENTS = "unifi_wireless_clients"
CONF_ALLOW_BANDWIDTH_SENSORS = "allow_bandwidth_sensors"
CONF_BLOCK_CLIENT = "block_client"
CONF_DETECTION_TIME = "detection_time"
CONF_TRACK_CLIENTS = "track_clients"
@ -23,6 +24,7 @@ CONF_DONT_TRACK_CLIENTS = "dont_track_clients"
CONF_DONT_TRACK_DEVICES = "dont_track_devices"
CONF_DONT_TRACK_WIRED_CLIENTS = "dont_track_wired_clients"
DEFAULT_ALLOW_BANDWIDTH_SENSORS = False
DEFAULT_BLOCK_CLIENTS = []
DEFAULT_TRACK_CLIENTS = True
DEFAULT_TRACK_DEVICES = True

15
homeassistant/components/unifi/controller.py
View File

@ -15,6 +15,7 @@ from homeassistant.helpers import aiohttp_client
from homeassistant.helpers.dispatcher import async_dispatcher_send
from .const import (
CONF_ALLOW_BANDWIDTH_SENSORS,
CONF_BLOCK_CLIENT,
CONF_CONTROLLER,
CONF_DETECTION_TIME,
@ -27,6 +28,7 @@ from .const import (
CONF_SITE_ID,
CONF_SSID_FILTER,
CONTROLLER_ID,
DEFAULT_ALLOW_BANDWIDTH_SENSORS,
DEFAULT_BLOCK_CLIENTS,
DEFAULT_TRACK_CLIENTS,
DEFAULT_TRACK_DEVICES,
@ -40,6 +42,8 @@ from .const import (
)
from .errors import AuthenticationRequired, CannotConnect
SUPPORTED_PLATFORMS = ["device_tracker", "sensor", "switch"]
class UniFiController:
"""Manages a single UniFi Controller."""
@ -76,6 +80,13 @@ class UniFiController:
"""Return the site user role of this controller."""
return self._site_role
@property
def option_allow_bandwidth_sensors(self):
"""Config entry option to allow bandwidth sensors."""
return self.config_entry.options.get(
CONF_ALLOW_BANDWIDTH_SENSORS, DEFAULT_ALLOW_BANDWIDTH_SENSORS
)
@property
def option_block_clients(self):
"""Config entry option with list of clients to control network access."""
@ -225,7 +236,7 @@ class UniFiController:
self.config_entry.add_update_listener(self.async_options_updated)
for platform in ["device_tracker", "switch"]:
for platform in SUPPORTED_PLATFORMS:
hass.async_create_task(
hass.config_entries.async_forward_entry_setup(
self.config_entry, platform
@ -294,7 +305,7 @@ class UniFiController:
if self.api is None:
return True
for platform in ["device_tracker", "switch"]:
for platform in SUPPORTED_PLATFORMS:
await self.hass.config_entries.async_forward_entry_unload(
self.config_entry, platform
)

168
homeassistant/components/unifi/sensor.py Normal file
View File

@ -0,0 +1,168 @@
"""Support for bandwidth sensors with UniFi clients."""
import logging
from homeassistant.components.unifi.config_flow import get_controller_from_config_entry
from homeassistant.core import callback
from homeassistant.helpers import entity_registry
from homeassistant.helpers.device_registry import CONNECTION_NETWORK_MAC
from homeassistant.helpers.dispatcher import async_dispatcher_connect
from homeassistant.helpers.entity import Entity
from homeassistant.helpers.entity_registry import DISABLED_CONFIG_ENTRY
LOGGER = logging.getLogger(__name__)
ATTR_RECEIVING = "receiving"
ATTR_TRANSMITTING = "transmitting"
async def async_setup_platform(hass, config, async_add_entities, discovery_info=None):
"""Sensor platform doesn't support configuration through configuration.yaml."""
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 = {}
registry = await entity_registry.async_get_registry(hass)
@callback
def update_controller():
"""Update the values of the controller."""
update_items(controller, async_add_entities, sensors)
async_dispatcher_connect(hass, controller.signal_update, update_controller)
@callback
def update_disable_on_entities():
"""Update the values of the controller."""
for entity in sensors.values():
disabled_by = None
if not entity.entity_registry_enabled_default and entity.enabled:
disabled_by = DISABLED_CONFIG_ENTRY
registry.async_update_entity(
entity.registry_entry.entity_id, disabled_by=disabled_by
)
async_dispatcher_connect(
hass, controller.signal_options_update, update_disable_on_entities
)
update_controller()
@callback
def update_items(controller, async_add_entities, sensors):
"""Update sensors from the controller."""
new_sensors = []
for client_id in controller.api.clients:
for direction, sensor_class in (
("rx", UniFiRxBandwidthSensor),
("tx", UniFiTxBandwidthSensor),
):
item_id = f"{direction}-{client_id}"
if item_id in sensors:
sensor = sensors[item_id]
if sensor.enabled:
sensor.async_schedule_update_ha_state()
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)
class UniFiBandwidthSensor(Entity):
"""UniFi Bandwidth sensor base class."""
def __init__(self, client, controller):
"""Set up client."""
self.client = client
self.controller = controller
self.is_wired = self.client.mac not in controller.wireless_clients
@property
def entity_registry_enabled_default(self):
"""Return if the entity should be enabled when first added to the entity registry."""
if self.controller.option_allow_bandwidth_sensors:
return True
return False
async def async_added_to_hass(self):
"""Client entity created."""
LOGGER.debug("New UniFi bandwidth sensor %s (%s)", self.name, self.client.mac)
async def async_update(self):
"""Synchronize state with controller.
Make sure to update self.is_wired if client is wireless, there is an issue when clients go offline that they get marked as wired.
"""
LOGGER.debug(
"Updating UniFi bandwidth sensor %s (%s)", self.entity_id, self.client.mac
)
await self.controller.request_update()
if self.is_wired and self.client.mac in self.controller.wireless_clients:
self.is_wired = False
@property
def available(self) -> bool:
"""Return if controller is available."""
return self.controller.available
@property
def device_info(self):
"""Return a device description for device registry."""
return {"connections": {(CONNECTION_NETWORK_MAC, self.client.mac)}}
class UniFiRxBandwidthSensor(UniFiBandwidthSensor):
"""Receiving bandwidth sensor."""
@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):
"""Return the name of the client."""
name = self.client.name or self.client.hostname
return f"{name} RX"
@property
def unique_id(self):
"""Return a unique identifier for this bandwidth sensor."""
return f"rx-{self.client.mac}"
class UniFiTxBandwidthSensor(UniFiBandwidthSensor):
"""Transmitting bandwidth sensor."""
@property
def state(self):
"""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}"

5
homeassistant/components/unifi/strings.json
View File

@ -35,6 +35,11 @@
"track_devices": "Track network devices (Ubiquiti devices)",
"track_wired_clients": "Include wired network clients"
}
},
"statistics_sensors": {
"data": {
"allow_bandwidth_sensors": "Create bandwidth usage sensors for network clients"
}
}
}
}

3
homeassistant/components/unifi/switch.py
View File

@ -14,7 +14,6 @@ LOGGER = logging.getLogger(__name__)
async def async_setup_platform(hass, config, async_add_entities, discovery_info=None):
"""Component doesn't support configuration through configuration.yaml."""
pass
async def async_setup_entry(hass, config_entry, async_add_entities):
@ -231,8 +230,6 @@ class UniFiPOEClientSwitch(UniFiClient, SwitchDevice, RestoreEntity):
"""Return the device state attributes."""
attributes = {
"power": self.port.poe_power,
"received": self.client.wired_rx_bytes / 1000000,
"sent": self.client.wired_tx_bytes / 1000000,
"switch": self.client.sw_mac,
"port": self.client.sw_port,
"poe_mode": self.poe_mode,

16
tests/components/unifi/test_controller.py
View File

@ -67,12 +67,18 @@ async def test_controller_setup():
assert await unifi_controller.async_setup() is True
assert unifi_controller.api is api
assert len(hass.config_entries.async_forward_entry_setup.mock_calls) == 2
assert len(hass.config_entries.async_forward_entry_setup.mock_calls) == len(
controller.SUPPORTED_PLATFORMS
)
assert hass.config_entries.async_forward_entry_setup.mock_calls[0][1] == (
entry,
"device_tracker",
)
assert hass.config_entries.async_forward_entry_setup.mock_calls[1][1] == (
entry,
"sensor",
)
assert hass.config_entries.async_forward_entry_setup.mock_calls[2][1] == (
entry,
"switch",
)
@ -214,12 +220,16 @@ async def test_reset_unloads_entry_if_setup():
with patch.object(controller, "get_controller", return_value=mock_coro(api)):
assert await unifi_controller.async_setup() is True
assert len(hass.config_entries.async_forward_entry_setup.mock_calls) == 2
assert len(hass.config_entries.async_forward_entry_setup.mock_calls) == len(
controller.SUPPORTED_PLATFORMS
)
hass.config_entries.async_forward_entry_unload.return_value = mock_coro(True)
assert await unifi_controller.async_reset()
assert len(hass.config_entries.async_forward_entry_unload.mock_calls) == 2
assert len(hass.config_entries.async_forward_entry_unload.mock_calls) == len(
controller.SUPPORTED_PLATFORMS
)
async def test_get_controller(hass):

207
tests/components/unifi/test_sensor.py Normal file
View File

@ -0,0 +1,207 @@
"""UniFi sensor platform tests."""
from collections import deque
from copy import deepcopy
from asynctest import patch
from homeassistant import config_entries
from homeassistant.components import unifi
from homeassistant.components.unifi.const import (
CONF_CONTROLLER,
CONF_SITE_ID,
CONTROLLER_ID as CONF_CONTROLLER_ID,
UNIFI_CONFIG,
UNIFI_WIRELESS_CLIENTS,
)
from homeassistant.setup import async_setup_component
from homeassistant.const import (
CONF_HOST,
CONF_PASSWORD,
CONF_PORT,
CONF_USERNAME,
CONF_VERIFY_SSL,
)
import homeassistant.components.sensor as sensor
CLIENTS = [
{
"hostname": "Wired client hostname",
"ip": "10.0.0.1",
"is_wired": True,
"last_seen": 1562600145,
"mac": "00:00:00:00:00:01",
"name": "Wired client name",
"oui": "Producer",
"sw_mac": "00:00:00:00:01:01",
"sw_port": 1,
"wired-rx_bytes": 1234000000,
"wired-tx_bytes": 5678000000,
},
{
"hostname": "Wireless client hostname",
"ip": "10.0.0.2",
"is_wired": False,
"last_seen": 1562600145,
"mac": "00:00:00:00:00:02",
"name": "Wireless client name",
"oui": "Producer",
"sw_mac": "00:00:00:00:01:01",
"sw_port": 2,
"rx_bytes": 1234000000,
"tx_bytes": 5678000000,
},
]
CONTROLLER_DATA = {
CONF_HOST: "mock-host",
CONF_USERNAME: "mock-user",
CONF_PASSWORD: "mock-pswd",
CONF_PORT: 1234,
CONF_SITE_ID: "mock-site",
CONF_VERIFY_SSL: False,
}
ENTRY_CONFIG = {CONF_CONTROLLER: CONTROLLER_DATA}
CONTROLLER_ID = CONF_CONTROLLER_ID.format(host="mock-host", site="mock-site")
SITES = {"Site name": {"desc": "Site name", "name": "mock-site", "role": "admin"}}
async def setup_unifi_integration(
hass,
config,
options,
sites,
clients_response,
devices_response,
clients_all_response,
):
"""Create the UniFi controller."""
hass.data[UNIFI_CONFIG] = []
hass.data[UNIFI_WIRELESS_CLIENTS] = unifi.UnifiWirelessClients(hass)
config_entry = config_entries.ConfigEntry(
version=1,
domain=unifi.DOMAIN,
title="Mock Title",
data=config,
source="test",
connection_class=config_entries.CONN_CLASS_LOCAL_POLL,
system_options={},
options=options,
entry_id=1,
)
mock_client_responses = deque()
mock_client_responses.append(clients_response)
mock_device_responses = deque()
mock_device_responses.append(devices_response)
mock_client_all_responses = deque()
mock_client_all_responses.append(clients_all_response)
mock_requests = []
async def mock_request(self, method, path, json=None):
mock_requests.append({"method": method, "path": path, "json": json})
if path == "s/{site}/stat/sta" and mock_client_responses:
return mock_client_responses.popleft()
if path == "s/{site}/stat/device" and mock_device_responses:
return mock_device_responses.popleft()
if path == "s/{site}/rest/user" and mock_client_all_responses:
return mock_client_all_responses.popleft()
return {}
with patch("aiounifi.Controller.login", return_value=True), patch(
"aiounifi.Controller.sites", return_value=sites
), patch("aiounifi.Controller.request", new=mock_request):
await unifi.async_setup_entry(hass, config_entry)
await hass.async_block_till_done()
hass.config_entries._entries.append(config_entry)
controller_id = unifi.get_controller_id_from_config_entry(config_entry)
controller = hass.data[unifi.DOMAIN][controller_id]
controller.mock_client_responses = mock_client_responses
controller.mock_device_responses = mock_device_responses
controller.mock_client_all_responses = mock_client_all_responses
controller.mock_requests = mock_requests
return controller
async def test_platform_manually_configured(hass):
"""Test that we do not discover anything or try to set up a controller."""
assert (
await async_setup_component(
hass, sensor.DOMAIN, {sensor.DOMAIN: {"platform": "unifi"}}
)
is True
)
assert unifi.DOMAIN not in hass.data
async def test_no_clients(hass):
"""Test the update_clients function when no clients are found."""
controller = await setup_unifi_integration(
hass,
ENTRY_CONFIG,
options={unifi.const.CONF_ALLOW_BANDWIDTH_SENSORS: True},
sites=SITES,
clients_response=[],
devices_response=[],
clients_all_response=[],
)
assert len(controller.mock_requests) == 3
assert len(hass.states.async_all()) == 2
async def test_switches(hass):
"""Test the update_items function with some clients."""
controller = await setup_unifi_integration(
hass,
ENTRY_CONFIG,
options={
unifi.const.CONF_ALLOW_BANDWIDTH_SENSORS: True,
unifi.const.CONF_TRACK_CLIENTS: False,
unifi.const.CONF_TRACK_DEVICES: False,
},
sites=SITES,
clients_response=CLIENTS,
devices_response=[],
clients_all_response=[],
)
assert len(controller.mock_requests) == 3
assert len(hass.states.async_all()) == 6
wired_client_rx = hass.states.get("sensor.wired_client_name_rx")
assert wired_client_rx.state == "1234.0"
wired_client_tx = hass.states.get("sensor.wired_client_name_tx")
assert wired_client_tx.state == "5678.0"
wireless_client_rx = hass.states.get("sensor.wireless_client_name_rx")
assert wireless_client_rx.state == "1234.0"
wireless_client_tx = hass.states.get("sensor.wireless_client_name_tx")
assert wireless_client_tx.state == "5678.0"
clients = deepcopy(CLIENTS)
clients[0]["is_wired"] = False
clients[1]["rx_bytes"] = 2345000000
clients[1]["tx_bytes"] = 6789000000
controller.mock_client_responses.append(clients)
await controller.async_update()
await hass.async_block_till_done()
wireless_client_rx = hass.states.get("sensor.wireless_client_name_rx")
assert wireless_client_rx.state == "2345.0"
wireless_client_tx = hass.states.get("sensor.wireless_client_name_tx")
assert wireless_client_tx.state == "6789.0"

4
tests/components/unifi/test_switch.py
View File

@ -264,7 +264,7 @@ async def setup_unifi_integration(
async def mock_request(self, method, path, json=None):
mock_requests.append({"method": method, "path": path, "json": json})
print(mock_requests, mock_client_responses, mock_device_responses)
if path == "s/{site}/stat/sta" and mock_client_responses:
return mock_client_responses.popleft()
if path == "s/{site}/stat/device" and mock_device_responses:
@ -386,8 +386,6 @@ async def test_switches(hass):
assert switch_1 is not None
assert switch_1.state == "on"
assert switch_1.attributes["power"] == "2.56"
assert switch_1.attributes["received"] == 1234
assert switch_1.attributes["sent"] == 5678
assert switch_1.attributes["switch"] == "00:00:00:00:01:01"
assert switch_1.attributes["port"] == 1
assert switch_1.attributes["poe_mode"] == "auto"