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