From b578a76efa96c24295741a03cf74c7f0b14c3de0 Mon Sep 17 00:00:00 2001 From: Robert Svensson Date: Thu, 23 Apr 2020 21:29:38 +0200 Subject: [PATCH] UniFi - Move some preloading of unavailable clients earlier in setup phase (#34599) Improve readability of setting up switch platform --- homeassistant/components/unifi/controller.py | 24 +++++++++ .../components/unifi/device_tracker.py | 20 ------- homeassistant/components/unifi/switch.py | 54 ++++++++----------- tests/components/unifi/test_switch.py | 4 +- 4 files changed, 48 insertions(+), 54 deletions(-) diff --git a/homeassistant/components/unifi/controller.py b/homeassistant/components/unifi/controller.py index 807e727777c..8ccb42794ec 100644 --- a/homeassistant/components/unifi/controller.py +++ b/homeassistant/components/unifi/controller.py @@ -265,6 +265,30 @@ class UniFiController: LOGGER.error("Unknown error connecting with UniFi controller: %s", err) return False + # Restore clients that is not a part of active clients list. + entity_registry = await self.hass.helpers.entity_registry.async_get_registry() + for entity in entity_registry.entities.values(): + if ( + entity.config_entry_id != self.config_entry.entry_id + or "-" not in entity.unique_id + ): + continue + + mac = "" + if entity.domain == TRACKER_DOMAIN: + mac, _ = entity.unique_id.split("-", 1) + elif entity.domain == SWITCH_DOMAIN: + _, mac = entity.unique_id.split("-", 1) + + if mac in self.api.clients or mac not in self.api.clients_all: + continue + + client = self.api.clients_all[mac] + self.api.clients.process_raw([client.raw]) + LOGGER.debug( + "Restore disconnected client %s (%s)", entity.entity_id, client.mac, + ) + wireless_clients = self.hass.data[UNIFI_WIRELESS_CLIENTS] self.wireless_clients = wireless_clients.get_data(self.config_entry) self.update_wireless_clients() diff --git a/homeassistant/components/unifi/device_tracker.py b/homeassistant/components/unifi/device_tracker.py index 4b72c520b31..7bc1bbd3197 100644 --- a/homeassistant/components/unifi/device_tracker.py +++ b/homeassistant/components/unifi/device_tracker.py @@ -47,26 +47,6 @@ async def async_setup_entry(hass, config_entry, async_add_entities): controller = hass.data[UNIFI_DOMAIN][config_entry.entry_id] controller.entities[DOMAIN] = {CLIENT_TRACKER: set(), DEVICE_TRACKER: set()} - # Restore clients that is not a part of active clients list. - entity_registry = await hass.helpers.entity_registry.async_get_registry() - for entity in entity_registry.entities.values(): - - if ( - entity.config_entry_id == config_entry.entry_id - and entity.domain == DOMAIN - and "-" in entity.unique_id - ): - - mac, _ = entity.unique_id.split("-", 1) - if mac in controller.api.clients or mac not in controller.api.clients_all: - continue - - client = controller.api.clients_all[mac] - controller.api.clients.process_raw([client.raw]) - LOGGER.debug( - "Restore disconnected client %s (%s)", entity.entity_id, client.mac, - ) - @callback def items_added(): """Update the values of the controller.""" diff --git a/homeassistant/components/unifi/switch.py b/homeassistant/components/unifi/switch.py index 73e1dd131bc..b562b8bd746 100644 --- a/homeassistant/components/unifi/switch.py +++ b/homeassistant/components/unifi/switch.py @@ -30,63 +30,53 @@ async def async_setup_entry(hass, config_entry, async_add_entities): if controller.site_role != "admin": return - switches_off = [] - - # Restore clients that is not a part of active clients list. + # Store previously known POE control entities in case their POE are turned off. + previously_known_poe_clients = [] entity_registry = await hass.helpers.entity_registry.async_get_registry() for entity in entity_registry.entities.values(): if ( - entity.config_entry_id == config_entry.entry_id - and entity.unique_id.startswith(f"{POE_SWITCH}-") + entity.config_entry_id != config_entry.entry_id + or not entity.unique_id.startswith(POE_SWITCH) ): + continue - _, mac = entity.unique_id.split("-", 1) + mac = entity.unique_id.replace(f"{POE_SWITCH}-", "") + if mac in controller.api.clients or mac in controller.api.clients_all: + previously_known_poe_clients.append(entity.unique_id) - if mac in controller.api.clients: - switches_off.append(entity.unique_id) - continue - - if mac in controller.api.clients_all: - client = controller.api.clients_all[mac] - controller.api.clients.process_raw([client.raw]) - switches_off.append(entity.unique_id) - continue + for mac in controller.option_block_clients: + if mac not in controller.api.clients and mac in controller.api.clients_all: + client = controller.api.clients_all[mac] + controller.api.clients.process_raw([client.raw]) @callback def items_added(): """Update the values of the controller.""" if controller.option_block_clients or controller.option_poe_clients: - add_entities(controller, async_add_entities, switches_off) + add_entities(controller, async_add_entities, previously_known_poe_clients) for signal in (controller.signal_update, controller.signal_options_update): controller.listeners.append(async_dispatcher_connect(hass, signal, items_added)) items_added() - switches_off.clear() + previously_known_poe_clients.clear() @callback -def add_entities(controller, async_add_entities, switches_off): +def add_entities(controller, async_add_entities, previously_known_poe_clients): """Add new switch entities from the controller.""" switches = [] for mac in controller.option_block_clients: - if mac in controller.entities[DOMAIN][BLOCK_SWITCH]: - continue - - client = None - - if mac in controller.api.clients: - client = controller.api.clients[mac] - - elif mac in controller.api.clients_all: - client = controller.api.clients_all[mac] - - if not client: + if ( + mac in controller.entities[DOMAIN][BLOCK_SWITCH] + or mac not in controller.api.clients + ): continue + client = controller.api.clients[mac] switches.append(UniFiBlockClientSwitch(client, controller)) if controller.option_poe_clients: @@ -101,7 +91,7 @@ def add_entities(controller, async_add_entities, switches_off): client = controller.api.clients[mac] - if poe_client_id not in switches_off and ( + if poe_client_id not in previously_known_poe_clients and ( mac in controller.wireless_clients or client.sw_mac not in devices or not devices[client.sw_mac].ports[client.sw_port].port_poe @@ -114,7 +104,7 @@ def add_entities(controller, async_add_entities, switches_off): multi_clients_on_port = False for client2 in controller.api.clients.values(): - if poe_client_id in switches_off: + if poe_client_id in previously_known_poe_clients: break if ( diff --git a/tests/components/unifi/test_switch.py b/tests/components/unifi/test_switch.py index 4dfd3bc4c9d..506fa14377c 100644 --- a/tests/components/unifi/test_switch.py +++ b/tests/components/unifi/test_switch.py @@ -383,7 +383,7 @@ async def test_option_block_clients(hass): options={CONF_BLOCK_CLIENT: [BLOCKED["mac"], UNBLOCKED["mac"]]}, ) await hass.async_block_till_done() - assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 2 + assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 1 # Remove the second switch again hass.config_entries.async_update_entry( @@ -397,7 +397,7 @@ async def test_option_block_clients(hass): controller.config_entry, options={CONF_BLOCK_CLIENT: [UNBLOCKED["mac"]]}, ) await hass.async_block_till_done() - assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 1 + assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 0 # Remove one hass.config_entries.async_update_entry(