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

@@ -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():
+        for entry in async_entries_for_config_entry(
-            if (
+            entity_registry, self.config_entry.entry_id
-                entity.config_entry_id != self.config_entry.entry_id
+        ):
-                or "-" not in entity.unique_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,
             )
 

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 (
+        if not entry.unique_id.startswith(POE_SWITCH):
-            entity.config_entry_id != config_entry.entry_id
-            or not entity.unique_id.startswith(POE_SWITCH)
-        ):
             continue
 
-        mac = entity.unique_id.replace(f"{POE_SWITCH}-", "")
+        mac = entry.unique_id.replace(f"{POE_SWITCH}-", "")
-        if mac in controller.api.clients or mac in controller.api.clients_all:
+        if mac not in controller.api.clients:
-            previously_known_poe_clients.append(entity.unique_id)
+            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(
+            add_poe_entities(controller, async_add_entities, clients, known_poe_clients)
-                controller, async_add_entities, clients, previously_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(
+def add_poe_entities(controller, async_add_entities, clients, known_poe_clients):
-    controller, async_add_entities, clients, previously_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 self.poe_mode:  # POE is enabled and client in a known state
-        if state is None:
             return
 
-        if self.poe_mode is None:
+        if (state := await self.async_get_last_state()) is None:
-            self.poe_mode = state.attributes["poe_mode"]
+            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
+            self.poe_mode is not None
-            or self.client.sw_mac
+            and self.controller.available
-            and (
+            and self.client.sw_port
-                self.controller.available
+            and self.client.sw_mac
-                and self.client.sw_mac in self.controller.api.devices
+            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]
-            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,
-            )
 
     async def options_updated(self) -> None:
         """Config entry options are updated, remove entity if option is disabled."""

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):
+async def test_restore_client_succeed(hass, aioclient_mock):
-    """Test the update_items function with some clients."""
+    """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"]}',
+        f'{POE_SWITCH}-{POE_CLIENT["mac"]}',
-        suggested_object_id=CLIENT_1["hostname"],
+        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"]}',
+        f'{POE_SWITCH}-{POE_CLIENT["mac"]}',
-        suggested_object_id=CLIENT_2["hostname"],
+        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],
+        clients_response=[],
-        devices_response=[DEVICE_1],
+        devices_response=[POE_DEVICE],
-        clients_all_response=[CLIENT_1],
+        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")
+    poe_client = hass.states.get("switch.poe_client")
-    assert device_1 is not None
+    assert poe_client.state == "unavailable"  # self.poe_mode is None