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
pull/47480/head
Robert Svensson 2021-03-05 22:09:05 +01:00 committed by GitHub
parent 02e723f206
commit 50d3aae418
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 200 additions and 68 deletions

View File

@ -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,
)

View File

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

View File

@ -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