From 50d3aae418cfb06f205eb7cac087a0be621bb6a4 Mon Sep 17 00:00:00 2001 From: Robert Svensson Date: Fri, 5 Mar 2021 22:09:05 +0100 Subject: [PATCH] Improve restoring UniFi POE entity state (#47148) * Improve restoring data and better handling when the restore data is empty Improve readability of some logic related to POE clients * There is no need to check clients_all in Switch platform * Add better tests when restoring state * Port except handling shouldn't be needed anymore * Walrusify get_last_state --- homeassistant/components/unifi/controller.py | 24 ++- homeassistant/components/unifi/switch.py | 73 ++++---- tests/components/unifi/test_switch.py | 171 +++++++++++++++++-- 3 files changed, 200 insertions(+), 68 deletions(-) diff --git a/homeassistant/components/unifi/controller.py b/homeassistant/components/unifi/controller.py index 7c95058e8e4..9096620f0ed 100644 --- a/homeassistant/components/unifi/controller.py +++ b/homeassistant/components/unifi/controller.py @@ -40,6 +40,7 @@ from homeassistant.core import callback from homeassistant.exceptions import ConfigEntryNotReady from homeassistant.helpers import aiohttp_client from homeassistant.helpers.dispatcher import async_dispatcher_send +from homeassistant.helpers.entity_registry import async_entries_for_config_entry from homeassistant.helpers.event import async_track_time_interval import homeassistant.util.dt as dt_util @@ -338,21 +339,18 @@ class UniFiController: self._site_role = description[0]["site_role"] - # Restore clients that is not a part of active clients list. + # Restore clients that are 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 - ): + for entry in async_entries_for_config_entry( + entity_registry, self.config_entry.entry_id + ): + if entry.domain == TRACKER_DOMAIN: + mac = entry.unique_id.split("-", 1)[0] + elif entry.domain == SWITCH_DOMAIN: + mac = entry.unique_id.split("-", 1)[1] + else: continue - mac = "" - if entity.domain == TRACKER_DOMAIN: - mac = entity.unique_id.split("-", 1)[0] - elif entity.domain == SWITCH_DOMAIN: - mac = entity.unique_id.split("-", 1)[1] - if mac in self.api.clients or mac not in self.api.clients_all: continue @@ -360,7 +358,7 @@ class UniFiController: self.api.clients.process_raw([client.raw]) LOGGER.debug( "Restore disconnected client %s (%s)", - entity.entity_id, + entry.entity_id, client.mac, ) diff --git a/homeassistant/components/unifi/switch.py b/homeassistant/components/unifi/switch.py index e596e0b1e2a..6c8b6ea35cd 100644 --- a/homeassistant/components/unifi/switch.py +++ b/homeassistant/components/unifi/switch.py @@ -18,6 +18,7 @@ from aiounifi.events import ( from homeassistant.components.switch import DOMAIN, SwitchEntity from homeassistant.core import callback from homeassistant.helpers.dispatcher import async_dispatcher_connect +from homeassistant.helpers.entity_registry import async_entries_for_config_entry from homeassistant.helpers.restore_state import RestoreEntity from .const import ATTR_MANUFACTURER, DOMAIN as UNIFI_DOMAIN @@ -50,19 +51,18 @@ async def async_setup_entry(hass, config_entry, async_add_entities): return # Store previously known POE control entities in case their POE are turned off. - previously_known_poe_clients = [] + known_poe_clients = [] entity_registry = await hass.helpers.entity_registry.async_get_registry() - for entity in entity_registry.entities.values(): + for entry in async_entries_for_config_entry(entity_registry, config_entry.entry_id): - if ( - entity.config_entry_id != config_entry.entry_id - or not entity.unique_id.startswith(POE_SWITCH) - ): + if not entry.unique_id.startswith(POE_SWITCH): continue - 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) + mac = entry.unique_id.replace(f"{POE_SWITCH}-", "") + if mac not in controller.api.clients: + continue + + known_poe_clients.append(mac) for mac in controller.option_block_clients: if mac not in controller.api.clients and mac in controller.api.clients_all: @@ -80,9 +80,7 @@ async def async_setup_entry(hass, config_entry, async_add_entities): add_block_entities(controller, async_add_entities, clients) if controller.option_poe_clients: - add_poe_entities( - controller, async_add_entities, clients, previously_known_poe_clients - ) + add_poe_entities(controller, async_add_entities, clients, known_poe_clients) if controller.option_dpi_restrictions: add_dpi_entities(controller, async_add_entities, dpi_groups) @@ -91,7 +89,7 @@ async def async_setup_entry(hass, config_entry, async_add_entities): controller.listeners.append(async_dispatcher_connect(hass, signal, items_added)) items_added() - previously_known_poe_clients.clear() + known_poe_clients.clear() @callback @@ -111,9 +109,7 @@ def add_block_entities(controller, async_add_entities, clients): @callback -def add_poe_entities( - controller, async_add_entities, clients, previously_known_poe_clients -): +def add_poe_entities(controller, async_add_entities, clients, known_poe_clients): """Add new switch entities from the controller.""" switches = [] @@ -123,10 +119,13 @@ def add_poe_entities( if mac in controller.entities[DOMAIN][POE_SWITCH]: continue - poe_client_id = f"{POE_SWITCH}-{mac}" client = controller.api.clients[mac] - if poe_client_id not in previously_known_poe_clients and ( + # Try to identify new clients powered by POE. + # Known POE clients have been created in previous HASS sessions. + # If port_poe is None the port does not support POE + # If poe_enable is False we can't know if a POE client is available for control. + if mac not in 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 @@ -139,7 +138,7 @@ def add_poe_entities( multi_clients_on_port = False for client2 in controller.api.clients.values(): - if poe_client_id in previously_known_poe_clients: + if mac in known_poe_clients: break if ( @@ -196,18 +195,19 @@ class UniFiPOEClientSwitch(UniFiClient, SwitchEntity, RestoreEntity): """Call when entity about to be added to Home Assistant.""" await super().async_added_to_hass() - state = await self.async_get_last_state() - if state is None: + if self.poe_mode: # POE is enabled and client in a known state return - if self.poe_mode is None: - self.poe_mode = state.attributes["poe_mode"] + if (state := await self.async_get_last_state()) is None: + return + + self.poe_mode = state.attributes.get("poe_mode") if not self.client.sw_mac: - self.client.raw["sw_mac"] = state.attributes["switch"] + self.client.raw["sw_mac"] = state.attributes.get("switch") if not self.client.sw_port: - self.client.raw["sw_port"] = state.attributes["port"] + self.client.raw["sw_port"] = state.attributes.get("port") @property def is_on(self): @@ -218,16 +218,15 @@ class UniFiPOEClientSwitch(UniFiClient, SwitchEntity, RestoreEntity): def available(self): """Return if switch is available. - Poe_mode None means its poe state is unknown. + Poe_mode None means its POE state is unknown. Sw_mac unavailable means restored client. """ return ( - self.poe_mode is None - or self.client.sw_mac - and ( - self.controller.available - and self.client.sw_mac in self.controller.api.devices - ) + self.poe_mode is not None + and self.controller.available + and self.client.sw_port + and self.client.sw_mac + and self.client.sw_mac in self.controller.api.devices ) async def async_turn_on(self, **kwargs): @@ -257,15 +256,7 @@ class UniFiPOEClientSwitch(UniFiClient, SwitchEntity, RestoreEntity): @property def port(self): """Shortcut to the switch port that client is connected to.""" - try: - return self.device.ports[self.client.sw_port] - except (AttributeError, KeyError, TypeError): - _LOGGER.warning( - "Entity %s reports faulty device %s or port %s", - self.entity_id, - self.client.sw_mac, - self.client.sw_port, - ) + return self.device.ports[self.client.sw_port] async def options_updated(self) -> None: """Config entry options are updated, remove entity if option is disabled.""" diff --git a/tests/components/unifi/test_switch.py b/tests/components/unifi/test_switch.py index a0dc8e984e1..44120a18ee3 100644 --- a/tests/components/unifi/test_switch.py +++ b/tests/components/unifi/test_switch.py @@ -1,9 +1,10 @@ """UniFi switch platform tests.""" from copy import deepcopy +from unittest.mock import patch from aiounifi.controller import MESSAGE_CLIENT_REMOVED, MESSAGE_EVENT -from homeassistant import config_entries +from homeassistant import config_entries, core from homeassistant.components.device_tracker import DOMAIN as TRACKER_DOMAIN from homeassistant.components.switch import DOMAIN as SWITCH_DOMAIN from homeassistant.components.unifi.const import ( @@ -726,8 +727,66 @@ async def test_ignore_multiple_poe_clients_on_same_port(hass, aioclient_mock): assert switch_2 is None -async def test_restoring_client(hass, aioclient_mock): - """Test the update_items function with some clients.""" +async def test_restore_client_succeed(hass, aioclient_mock): + """Test that RestoreEntity works as expected.""" + POE_DEVICE = { + "device_id": "12345", + "ip": "1.0.1.1", + "mac": "00:00:00:00:01:01", + "last_seen": 1562600145, + "model": "US16P150", + "name": "POE Switch", + "port_overrides": [ + { + "poe_mode": "off", + "port_idx": 1, + "portconf_id": "5f3edd2aba4cc806a19f2db2", + } + ], + "port_table": [ + { + "media": "GE", + "name": "Port 1", + "op_mode": "switch", + "poe_caps": 7, + "poe_class": "Unknown", + "poe_current": "0.00", + "poe_enable": False, + "poe_good": False, + "poe_mode": "off", + "poe_power": "0.00", + "poe_voltage": "0.00", + "port_idx": 1, + "port_poe": True, + "portconf_id": "5f3edd2aba4cc806a19f2db2", + "up": False, + }, + ], + "state": 1, + "type": "usw", + "version": "4.0.42.10433", + } + POE_CLIENT = { + "hostname": "poe_client", + "ip": "1.0.0.1", + "is_wired": True, + "last_seen": 1562600145, + "mac": "00:00:00:00:00:01", + "name": "POE Client", + "oui": "Producer", + } + + fake_state = core.State( + "switch.poe_client", + "off", + { + "power": "0.00", + "switch": POE_DEVICE["mac"], + "port": 1, + "poe_mode": "auto", + }, + ) + config_entry = config_entries.ConfigEntry( version=1, domain=UNIFI_DOMAIN, @@ -744,15 +803,100 @@ async def test_restoring_client(hass, aioclient_mock): registry.async_get_or_create( SWITCH_DOMAIN, UNIFI_DOMAIN, - f'{POE_SWITCH}-{CLIENT_1["mac"]}', - suggested_object_id=CLIENT_1["hostname"], + f'{POE_SWITCH}-{POE_CLIENT["mac"]}', + suggested_object_id=POE_CLIENT["hostname"], config_entry=config_entry, ) + + with patch( + "homeassistant.helpers.restore_state.RestoreEntity.async_get_last_state", + return_value=fake_state, + ): + await setup_unifi_integration( + hass, + aioclient_mock, + options={ + CONF_TRACK_CLIENTS: False, + CONF_TRACK_DEVICES: False, + }, + clients_response=[], + devices_response=[POE_DEVICE], + clients_all_response=[POE_CLIENT], + ) + + assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 1 + + poe_client = hass.states.get("switch.poe_client") + assert poe_client.state == "off" + + +async def test_restore_client_no_old_state(hass, aioclient_mock): + """Test that RestoreEntity without old state makes entity unavailable.""" + POE_DEVICE = { + "device_id": "12345", + "ip": "1.0.1.1", + "mac": "00:00:00:00:01:01", + "last_seen": 1562600145, + "model": "US16P150", + "name": "POE Switch", + "port_overrides": [ + { + "poe_mode": "off", + "port_idx": 1, + "portconf_id": "5f3edd2aba4cc806a19f2db2", + } + ], + "port_table": [ + { + "media": "GE", + "name": "Port 1", + "op_mode": "switch", + "poe_caps": 7, + "poe_class": "Unknown", + "poe_current": "0.00", + "poe_enable": False, + "poe_good": False, + "poe_mode": "off", + "poe_power": "0.00", + "poe_voltage": "0.00", + "port_idx": 1, + "port_poe": True, + "portconf_id": "5f3edd2aba4cc806a19f2db2", + "up": False, + }, + ], + "state": 1, + "type": "usw", + "version": "4.0.42.10433", + } + POE_CLIENT = { + "hostname": "poe_client", + "ip": "1.0.0.1", + "is_wired": True, + "last_seen": 1562600145, + "mac": "00:00:00:00:00:01", + "name": "POE Client", + "oui": "Producer", + } + + config_entry = config_entries.ConfigEntry( + version=1, + domain=UNIFI_DOMAIN, + title="Mock Title", + data=ENTRY_CONFIG, + source="test", + connection_class=config_entries.CONN_CLASS_LOCAL_POLL, + system_options={}, + options={}, + entry_id=1, + ) + + registry = await entity_registry.async_get_registry(hass) registry.async_get_or_create( SWITCH_DOMAIN, UNIFI_DOMAIN, - f'{POE_SWITCH}-{CLIENT_2["mac"]}', - suggested_object_id=CLIENT_2["hostname"], + f'{POE_SWITCH}-{POE_CLIENT["mac"]}', + suggested_object_id=POE_CLIENT["hostname"], config_entry=config_entry, ) @@ -760,16 +904,15 @@ async def test_restoring_client(hass, aioclient_mock): hass, aioclient_mock, options={ - CONF_BLOCK_CLIENT: ["random mac"], CONF_TRACK_CLIENTS: False, CONF_TRACK_DEVICES: False, }, - clients_response=[CLIENT_2], - devices_response=[DEVICE_1], - clients_all_response=[CLIENT_1], + clients_response=[], + devices_response=[POE_DEVICE], + clients_all_response=[POE_CLIENT], ) - assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 2 + assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 1 - device_1 = hass.states.get("switch.client_1") - assert device_1 is not None + poe_client = hass.states.get("switch.poe_client") + assert poe_client.state == "unavailable" # self.poe_mode is None