UniFi - Support automatic removal of clients (#34307)
This commit is contained in:
parent
abae48c287
commit
374fe47809
11 changed files with 227 additions and 46 deletions
|
@ -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."""
|
||||
|
|
|
@ -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."""
|
||||
|
||||
|
|
|
@ -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"
|
||||
}
|
||||
}
|
||||
|
|
|
@ -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."""
|
||||
|
||||
|
|
|
@ -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."""
|
||||
|
||||
|
|
|
@ -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
|
||||
|
|
|
@ -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
|
||||
|
|
|
@ -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()
|
||||
|
|
|
@ -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(
|
||||
|
|
|
@ -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
|
||||
|
|
|
@ -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(
|
||||
|
|
Loading…
Add table
Reference in a new issue