diff --git a/homeassistant/components/unifi/controller.py b/homeassistant/components/unifi/controller.py index 864d131d287..3f4944fafa9 100644 --- a/homeassistant/components/unifi/controller.py +++ b/homeassistant/components/unifi/controller.py @@ -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.""" diff --git a/homeassistant/components/unifi/device_tracker.py b/homeassistant/components/unifi/device_tracker.py index c5aa74706a1..3f2eaf698b9 100644 --- a/homeassistant/components/unifi/device_tracker.py +++ b/homeassistant/components/unifi/device_tracker.py @@ -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.""" diff --git a/homeassistant/components/unifi/manifest.json b/homeassistant/components/unifi/manifest.json index a58bcd6fa7a..0a5ba84cdb3 100644 --- a/homeassistant/components/unifi/manifest.json +++ b/homeassistant/components/unifi/manifest.json @@ -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" -} \ No newline at end of file +} diff --git a/homeassistant/components/unifi/sensor.py b/homeassistant/components/unifi/sensor.py index 2e82ecb4f6f..2777d21cac7 100644 --- a/homeassistant/components/unifi/sensor.py +++ b/homeassistant/components/unifi/sensor.py @@ -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.""" diff --git a/homeassistant/components/unifi/switch.py b/homeassistant/components/unifi/switch.py index 0547e35f064..a0b7d865a1b 100644 --- a/homeassistant/components/unifi/switch.py +++ b/homeassistant/components/unifi/switch.py @@ -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.""" diff --git a/requirements_all.txt b/requirements_all.txt index 7cd4accf59e..ce58b1780c8 100644 --- a/requirements_all.txt +++ b/requirements_all.txt @@ -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 diff --git a/requirements_test_all.txt b/requirements_test_all.txt index 92351e40108..789c23f21de 100644 --- a/requirements_test_all.txt +++ b/requirements_test_all.txt @@ -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 diff --git a/tests/components/unifi/test_controller.py b/tests/components/unifi/test_controller.py index ad334f848ba..d78947f3134 100644 --- a/tests/components/unifi/test_controller.py +++ b/tests/components/unifi/test_controller.py @@ -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() diff --git a/tests/components/unifi/test_device_tracker.py b/tests/components/unifi/test_device_tracker.py index cfb4637a6c4..84085f8711a 100644 --- a/tests/components/unifi/test_device_tracker.py +++ b/tests/components/unifi/test_device_tracker.py @@ -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( diff --git a/tests/components/unifi/test_sensor.py b/tests/components/unifi/test_sensor.py index 91531a9ee38..3c801474235 100644 --- a/tests/components/unifi/test_sensor.py +++ b/tests/components/unifi/test_sensor.py @@ -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 diff --git a/tests/components/unifi/test_switch.py b/tests/components/unifi/test_switch.py index cc777b8ad7f..5ea76472739 100644 --- a/tests/components/unifi/test_switch.py +++ b/tests/components/unifi/test_switch.py @@ -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(