UniFi POE control restore clients (#25558)

* Restore POE controls on restart

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."""

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