From dc722adbb53edc1f7b1b7bc5d5f0d7098cc7ab97 Mon Sep 17 00:00:00 2001 From: Robert Svensson Date: Mon, 29 Jul 2019 19:48:38 +0200 Subject: [PATCH] UniFi POE control restore clients (#25558) * Restore POE controls on restart --- homeassistant/components/unifi/switch.py | 73 ++++++++++++++++++++---- tests/components/unifi/test_switch.py | 32 ++++++++++- 2 files changed, 91 insertions(+), 14 deletions(-) diff --git a/homeassistant/components/unifi/switch.py b/homeassistant/components/unifi/switch.py index 3d3ccc95563..610f178c1a4 100644 --- a/homeassistant/components/unifi/switch.py +++ b/homeassistant/components/unifi/switch.py @@ -5,8 +5,10 @@ from homeassistant.components import unifi from homeassistant.components.switch import SwitchDevice from homeassistant.const import CONF_HOST from homeassistant.core import callback +from homeassistant.helpers import entity_registry from homeassistant.helpers.device_registry import CONNECTION_NETWORK_MAC from homeassistant.helpers.dispatcher import async_dispatcher_connect +from homeassistant.helpers.restore_state import RestoreEntity from .const import CONF_CONTROLLER, CONF_SITE_ID, CONTROLLER_ID @@ -34,19 +36,39 @@ async def async_setup_entry(hass, config_entry, async_add_entities): return switches = {} + switches_off = [] + + registry = await entity_registry.async_get_registry(hass) + + # Restore clients that is not a part of active clients list. + for entity in registry.entities.values(): + + if entity.config_entry_id == config_entry.entry_id and \ + entity.unique_id.startswith('poe-'): + + _, 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]) + switches_off.append(entity.unique_id) @callback def update_controller(): """Update the values of the controller.""" - update_items(controller, async_add_entities, switches) + update_items(controller, async_add_entities, switches, switches_off) async_dispatcher_connect(hass, controller.event_update, update_controller) update_controller() + switches_off.clear() @callback -def update_items(controller, async_add_entities, switches): +def update_items(controller, async_add_entities, switches, switches_off): """Update POE port state from the controller.""" new_switches = [] devices = controller.api.devices @@ -85,17 +107,23 @@ def update_items(controller, async_add_entities, switches): continue client = controller.api.clients[client_id] + + if poe_client_id in switches_off: + pass # Network device with active POE - if not client.is_wired or client.sw_mac not in devices or \ - not devices[client.sw_mac].ports[client.sw_port].port_poe or \ - not devices[client.sw_mac].ports[client.sw_port].poe_enable or \ - controller.mac == client.mac: + elif not client.is_wired or client.sw_mac not in devices or \ + not devices[client.sw_mac].ports[client.sw_port].port_poe or \ + controller.mac == client.mac: continue # Multiple POE-devices on same port means non UniFi POE driven switch multi_clients_on_port = False for client2 in controller.api.clients.values(): - if client.mac != client2.mac and \ + + if poe_client_id in switches_off: + break + + if client2.is_wired and client.mac != client2.mac and \ client.sw_mac == client2.sw_mac and \ client.sw_port == client2.sw_port: multi_clients_on_port = True @@ -138,16 +166,32 @@ class UniFiClient: } -class UniFiPOEClientSwitch(UniFiClient, SwitchDevice): +class UniFiPOEClientSwitch(UniFiClient, SwitchDevice, RestoreEntity): """Representation of a client that uses POE.""" def __init__(self, client, controller): """Set up POE switch.""" super().__init__(client, controller) self.poe_mode = None - if self.port.poe_mode != 'off': + if self.client.sw_port and self.port.poe_mode != 'off': self.poe_mode = self.port.poe_mode + async def async_added_to_hass(self): + """Call when entity about to be added to Home Assistant.""" + state = await self.async_get_last_state() + + if state is None: + return + + if self.poe_mode is None: + self.poe_mode = state.attributes['poe_mode'] + + if not self.client.sw_mac: + self.client.raw['sw_mac'] = state.attributes['switch'] + + if not self.client.sw_port: + self.client.raw['sw_port'] = state.attributes['port'] + @property def unique_id(self): """Return a unique identifier for this switch.""" @@ -160,9 +204,14 @@ class UniFiPOEClientSwitch(UniFiClient, SwitchDevice): @property def available(self): - """Return if switch is available.""" - return self.controller.available or \ - self.client.sw_mac in self.controller.api.devices + """Return if switch is available. + + 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 or + self.client.sw_mac in self.controller.api.devices) async def async_turn_on(self, **kwargs): """Enable POE for client.""" diff --git a/tests/components/unifi/test_switch.py b/tests/components/unifi/test_switch.py index f467adad9a2..2d9ea20aca2 100644 --- a/tests/components/unifi/test_switch.py +++ b/tests/components/unifi/test_switch.py @@ -14,6 +14,7 @@ from homeassistant import config_entries from homeassistant.components import unifi from homeassistant.components.unifi.const import ( CONF_CONTROLLER, CONF_SITE_ID, UNIFI_CONFIG) +from homeassistant.helpers import entity_registry from homeassistant.setup import async_setup_component from homeassistant.const import ( CONF_HOST, CONF_PASSWORD, CONF_PORT, CONF_USERNAME, CONF_VERIFY_SSL) @@ -249,7 +250,7 @@ async def setup_controller(hass, mock_controller): hass.data[unifi.DOMAIN] = {CONTROLLER_ID: mock_controller} config_entry = config_entries.ConfigEntry( 1, unifi.DOMAIN, 'Mock Title', ENTRY_CONFIG, 'test', - config_entries.CONN_CLASS_LOCAL_POLL) + config_entries.CONN_CLASS_LOCAL_POLL, entry_id=1) mock_controller.config_entry = config_entry await mock_controller.async_update() @@ -312,7 +313,7 @@ async def test_switches(hass, mock_controller): await setup_controller(hass, mock_controller) assert len(mock_controller.mock_requests) == 3 - assert len(hass.states.async_all()) == 4 + assert len(hass.states.async_all()) == 5 switch_1 = hass.states.get('switch.poe_client_1') assert switch_1 is not None @@ -450,3 +451,30 @@ async def test_ignore_multiple_poe_clients_on_same_port(hass, mock_controller): switch_2 = hass.states.get('switch.poe_client_2') assert switch_1 is None assert switch_2 is None + + +async def test_restoring_client(hass, mock_controller): + """Test the update_items function with some clients.""" + mock_controller.mock_client_responses.append([CLIENT_2]) + mock_controller.mock_device_responses.append([DEVICE_1]) + mock_controller.mock_client_all_responses.append([CLIENT_1]) + mock_controller.unifi_config = { + unifi.CONF_BLOCK_CLIENT: ['random mac'] + } + + registry = await entity_registry.async_get_registry(hass) + registry.async_get_or_create( + switch.DOMAIN, unifi.DOMAIN, + 'poe-{}'.format(CLIENT_1['mac']), + suggested_object_id=CLIENT_1['hostname'], config_entry_id=1) + registry.async_get_or_create( + switch.DOMAIN, unifi.DOMAIN, + 'poe-{}'.format(CLIENT_2['mac']), + suggested_object_id=CLIENT_2['hostname'], config_entry_id=1) + + await setup_controller(hass, mock_controller) + assert len(mock_controller.mock_requests) == 3 + assert len(hass.states.async_all()) == 3 + + device_1 = hass.states.get('switch.client_1') + assert device_1 is not None