UniFi - Support automatic removal of clients (#34307)

This commit is contained in:
Robert Svensson 2020-04-17 00:08:53 +02:00 committed by GitHub
parent abae48c287
commit 374fe47809
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
11 changed files with 227 additions and 46 deletions

View file

@ -5,7 +5,14 @@ import ssl
from aiohttp import CookieJar
import aiounifi
from aiounifi.controller import SIGNAL_CONNECTION_STATE
from aiounifi.controller import (
DATA_CLIENT,
DATA_CLIENT_REMOVED,
DATA_DEVICE,
DATA_EVENT,
SIGNAL_CONNECTION_STATE,
SIGNAL_DATA,
)
from aiounifi.events import WIRELESS_CLIENT_CONNECTED, WIRELESS_GUEST_CONNECTED
from aiounifi.websocket import STATE_DISCONNECTED, STATE_RUNNING
import async_timeout
@ -161,16 +168,23 @@ class UniFiController:
if not self.available:
self.hass.loop.call_later(RETRY_TIMER, self.reconnect)
elif signal == "new_data" and data:
if "event" in data:
if data["event"].event in (
elif signal == SIGNAL_DATA and data:
if DATA_EVENT in data:
if data[DATA_EVENT].event in (
WIRELESS_CLIENT_CONNECTED,
WIRELESS_GUEST_CONNECTED,
):
self.update_wireless_clients()
elif "clients" in data or "devices" in data:
elif DATA_CLIENT in data or DATA_DEVICE in data:
async_dispatcher_send(self.hass, self.signal_update)
elif DATA_CLIENT_REMOVED in data:
async_dispatcher_send(
self.hass, self.signal_remove, data[DATA_CLIENT_REMOVED]
)
@property
def signal_reachable(self) -> str:
"""Integration specific event to signal a change in connection status."""
@ -181,6 +195,11 @@ class UniFiController:
"""Event specific per UniFi entry to signal new data."""
return f"unifi-update-{self.controller_id}"
@property
def signal_remove(self):
"""Event specific per UniFi entry to signal removal of entities."""
return f"unifi-remove-{self.controller_id}"
@property
def signal_options_update(self):
"""Event specific per UniFi entry to signal new options."""

View file

@ -49,10 +49,10 @@ async def async_setup_entry(hass, config_entry, async_add_entities):
option_track_wired_clients = controller.option_track_wired_clients
option_ssid_filter = controller.option_ssid_filter
registry = await hass.helpers.entity_registry.async_get_registry()
entity_registry = await hass.helpers.entity_registry.async_get_registry()
# Restore clients that is not a part of active clients list.
for entity in registry.entities.values():
for entity in entity_registry.entities.values():
if (
entity.config_entry_id == config_entry.entry_id
@ -69,7 +69,7 @@ async def async_setup_entry(hass, config_entry, async_add_entities):
controller.api.clients.process_raw([client.raw])
@callback
def update_controller():
def items_added():
"""Update the values of the controller."""
nonlocal option_track_clients
nonlocal option_track_devices
@ -80,7 +80,16 @@ async def async_setup_entry(hass, config_entry, async_add_entities):
add_entities(controller, async_add_entities, tracked)
controller.listeners.append(
async_dispatcher_connect(hass, controller.signal_update, update_controller)
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, tracked, entity_registry)
controller.listeners.append(
async_dispatcher_connect(hass, controller.signal_remove, items_removed)
)
@callback
@ -136,16 +145,10 @@ async def async_setup_entry(hass, config_entry, async_add_entities):
option_track_devices = controller.option_track_devices
option_track_wired_clients = controller.option_track_wired_clients
for mac in remove:
entity = tracked.pop(mac)
if registry.async_is_registered(entity.entity_id):
registry.async_remove(entity.entity_id)
hass.async_create_task(entity.async_remove())
remove_entities(controller, remove, tracked, entity_registry)
if update:
update_controller()
items_added()
controller.listeners.append(
async_dispatcher_connect(
@ -153,7 +156,7 @@ async def async_setup_entry(hass, config_entry, async_add_entities):
)
)
update_controller()
items_added()
@callback
@ -193,6 +196,18 @@ def add_entities(controller, async_add_entities, tracked):
async_add_entities(new_tracked)
@callback
def remove_entities(controller, mac_addresses, tracked, entity_registry):
"""Remove select tracked entities."""
for mac in mac_addresses:
if mac not in tracked:
continue
entity = tracked.pop(mac)
controller.hass.async_create_task(entity.async_remove())
class UniFiClientTracker(UniFiClient, ScannerEntity):
"""Representation of a network client."""

View file

@ -3,11 +3,7 @@
"name": "Ubiquiti UniFi",
"config_flow": true,
"documentation": "https://www.home-assistant.io/integrations/unifi",
"requirements": [
"aiounifi==17"
],
"codeowners": [
"@kane610"
],
"requirements": ["aiounifi==18"],
"codeowners": ["@kane610"],
"quality_scale": "platinum"
}
}

View file

@ -25,7 +25,7 @@ async def async_setup_entry(hass, config_entry, async_add_entities):
entity_registry = await hass.helpers.entity_registry.async_get_registry()
@callback
def update_controller():
def items_added():
"""Update the values of the controller."""
nonlocal option_allow_bandwidth_sensors
@ -35,7 +35,16 @@ async def async_setup_entry(hass, config_entry, async_add_entities):
add_entities(controller, async_add_entities, sensors)
controller.listeners.append(
async_dispatcher_connect(hass, controller.signal_update, update_controller)
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
@ -47,14 +56,10 @@ async def async_setup_entry(hass, config_entry, async_add_entities):
option_allow_bandwidth_sensors = controller.option_allow_bandwidth_sensors
if option_allow_bandwidth_sensors:
update_controller()
items_added()
else:
for sensor in sensors.values():
if entity_registry.async_is_registered(sensor.entity_id):
entity_registry.async_remove(sensor.entity_id)
hass.async_create_task(sensor.async_remove())
sensors.clear()
@ -65,7 +70,7 @@ async def async_setup_entry(hass, config_entry, async_add_entities):
)
)
update_controller()
items_added()
@callback
@ -92,6 +97,21 @@ def add_entities(controller, async_add_entities, sensors):
async_add_entities(new_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."""

View file

@ -55,12 +55,21 @@ async def async_setup_entry(hass, config_entry, async_add_entities):
continue
@callback
def update_controller():
def items_added():
"""Update the values of the controller."""
add_entities(controller, async_add_entities, switches, switches_off)
controller.listeners.append(
async_dispatcher_connect(hass, controller.signal_update, update_controller)
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, switches, entity_registry)
controller.listeners.append(
async_dispatcher_connect(hass, controller.signal_remove, items_removed)
)
@callback
@ -96,14 +105,10 @@ async def async_setup_entry(hass, config_entry, async_add_entities):
for client_id in remove:
entity = switches.pop(client_id)
if entity_registry.async_is_registered(entity.entity_id):
entity_registry.async_remove(entity.entity_id)
hass.async_create_task(entity.async_remove())
if len(update) != len(option_block_clients):
update_controller()
items_added()
controller.listeners.append(
async_dispatcher_connect(
@ -111,7 +116,7 @@ async def async_setup_entry(hass, config_entry, async_add_entities):
)
)
update_controller()
items_added()
switches_off.clear()
@ -189,6 +194,21 @@ def add_entities(controller, async_add_entities, switches, switches_off):
async_add_entities(new_switches)
@callback
def remove_entities(controller, mac_addresses, switches, entity_registry):
"""Remove select switch entities."""
for mac in mac_addresses:
for switch_type in ("block", "poe"):
item_id = f"{switch_type}-{mac}"
if item_id not in switches:
continue
entity = switches.pop(item_id)
controller.hass.async_create_task(entity.async_remove())
class UniFiPOEClientSwitch(UniFiClient, SwitchDevice, RestoreEntity):
"""Representation of a client that uses POE."""

View file

@ -211,7 +211,7 @@ aiopylgtv==0.3.3
aioswitcher==1.1.0
# homeassistant.components.unifi
aiounifi==17
aiounifi==18
# homeassistant.components.wwlln
aiowwlln==2.0.2

View file

@ -94,7 +94,7 @@ aiopylgtv==0.3.3
aioswitcher==1.1.0
# homeassistant.components.unifi
aiounifi==17
aiounifi==18
# homeassistant.components.wwlln
aiowwlln==2.0.2

View file

@ -177,6 +177,7 @@ async def test_controller_setup(hass):
assert controller.mac is None
assert controller.signal_update == "unifi-update-1.2.3.4-site_id"
assert controller.signal_remove == "unifi-remove-1.2.3.4-site_id"
assert controller.signal_options_update == "unifi-options-1.2.3.4-site_id"
@ -206,7 +207,7 @@ async def test_reset_after_successful_setup(hass):
"""Calling reset when the entry has been setup."""
controller = await setup_unifi_integration(hass)
assert len(controller.listeners) == 6
assert len(controller.listeners) == 9
result = await controller.async_reset()
await hass.async_block_till_done()

View file

@ -2,8 +2,8 @@
from copy import copy
from datetime import timedelta
from aiounifi.controller import SIGNAL_CONNECTION_STATE
from aiounifi.websocket import STATE_DISCONNECTED, STATE_RUNNING
from aiounifi.controller import MESSAGE_CLIENT_REMOVED, SIGNAL_CONNECTION_STATE
from aiounifi.websocket import SIGNAL_DATA, STATE_DISCONNECTED, STATE_RUNNING
from asynctest import patch
from homeassistant import config_entries
@ -180,6 +180,35 @@ async def test_tracked_devices(hass):
assert device_1.state == STATE_UNAVAILABLE
async def test_remove_clients(hass):
"""Test the remove_items function with some clients."""
controller = await setup_unifi_integration(
hass, clients_response=[CLIENT_1, CLIENT_2]
)
assert len(hass.states.async_entity_ids("device_tracker")) == 2
client_1 = hass.states.get("device_tracker.client_1")
assert client_1 is not None
wired_client = hass.states.get("device_tracker.wired_client")
assert wired_client is not None
controller.api.websocket._data = {
"meta": {"message": MESSAGE_CLIENT_REMOVED},
"data": [CLIENT_1],
}
controller.api.session_handler(SIGNAL_DATA)
await hass.async_block_till_done()
assert len(hass.states.async_entity_ids("device_tracker")) == 1
client_1 = hass.states.get("device_tracker.client_1")
assert client_1 is None
wired_client = hass.states.get("device_tracker.wired_client")
assert wired_client is not None
async def test_controller_state_change(hass):
"""Verify entities state reflect on controller becoming unavailable."""
controller = await setup_unifi_integration(

View file

@ -1,6 +1,9 @@
"""UniFi sensor platform tests."""
from copy import deepcopy
from aiounifi.controller import MESSAGE_CLIENT_REMOVED
from aiounifi.websocket import SIGNAL_DATA
from homeassistant.components import unifi
import homeassistant.components.sensor as sensor
from homeassistant.setup import async_setup_component
@ -123,3 +126,44 @@ async def test_sensors(hass):
wireless_client_tx = hass.states.get("sensor.wireless_client_name_tx")
assert wireless_client_tx.state == "6789.0"
async def test_remove_sensors(hass):
"""Test the remove_items function with some clients."""
controller = await setup_unifi_integration(
hass,
options={unifi.const.CONF_ALLOW_BANDWIDTH_SENSORS: True},
clients_response=CLIENTS,
)
assert len(hass.states.async_entity_ids("sensor")) == 4
assert len(hass.states.async_entity_ids("device_tracker")) == 2
wired_client_rx = hass.states.get("sensor.wired_client_name_rx")
assert wired_client_rx is not None
wired_client_tx = hass.states.get("sensor.wired_client_name_tx")
assert wired_client_tx is not None
wireless_client_rx = hass.states.get("sensor.wireless_client_name_rx")
assert wireless_client_rx is not None
wireless_client_tx = hass.states.get("sensor.wireless_client_name_tx")
assert wireless_client_tx is not None
controller.api.websocket._data = {
"meta": {"message": MESSAGE_CLIENT_REMOVED},
"data": [CLIENTS[0]],
}
controller.api.session_handler(SIGNAL_DATA)
await hass.async_block_till_done()
assert len(hass.states.async_entity_ids("sensor")) == 2
assert len(hass.states.async_entity_ids("device_tracker")) == 1
wired_client_rx = hass.states.get("sensor.wired_client_name_rx")
assert wired_client_rx is None
wired_client_tx = hass.states.get("sensor.wired_client_name_tx")
assert wired_client_tx is None
wireless_client_rx = hass.states.get("sensor.wireless_client_name_rx")
assert wireless_client_rx is not None
wireless_client_tx = hass.states.get("sensor.wireless_client_name_tx")
assert wireless_client_tx is not None

View file

@ -1,6 +1,9 @@
"""UniFi POE control platform tests."""
from copy import deepcopy
from aiounifi.controller import MESSAGE_CLIENT_REMOVED
from aiounifi.websocket import SIGNAL_DATA
from homeassistant import config_entries
from homeassistant.components import unifi
import homeassistant.components.switch as switch
@ -173,6 +176,7 @@ BLOCKED = {
"ip": "10.0.0.1",
"is_guest": False,
"is_wired": False,
"last_seen": 1562600145,
"mac": "00:00:00:00:01:01",
"name": "Block Client 1",
"noted": True,
@ -184,6 +188,7 @@ UNBLOCKED = {
"ip": "10.0.0.2",
"is_guest": False,
"is_wired": True,
"last_seen": 1562600145,
"mac": "00:00:00:00:01:02",
"name": "Block Client 2",
"noted": True,
@ -300,6 +305,38 @@ async def test_switches(hass):
}
async def test_remove_switches(hass):
"""Test the update_items function with some clients."""
controller = await setup_unifi_integration(
hass,
options={CONF_BLOCK_CLIENT: [UNBLOCKED["mac"]]},
clients_response=[CLIENT_1, UNBLOCKED],
devices_response=[DEVICE_1],
)
assert len(hass.states.async_entity_ids("switch")) == 2
poe_switch = hass.states.get("switch.poe_client_1")
assert poe_switch is not None
block_switch = hass.states.get("switch.block_client_2")
assert block_switch is not None
controller.api.websocket._data = {
"meta": {"message": MESSAGE_CLIENT_REMOVED},
"data": [CLIENT_1, UNBLOCKED],
}
controller.api.session_handler(SIGNAL_DATA)
await hass.async_block_till_done()
assert len(hass.states.async_entity_ids("switch")) == 0
poe_switch = hass.states.get("switch.poe_client_1")
assert poe_switch is None
block_switch = hass.states.get("switch.block_client_2")
assert block_switch is None
async def test_new_client_discovered_on_block_control(hass):
"""Test if 2nd update has a new client."""
controller = await setup_unifi_integration(