UniFi POE control restore clients (#25558)

* Restore POE controls on restart
This commit is contained in:
Robert Svensson 2019-07-29 19:48:38 +02:00 committed by GitHub
parent 2e300aec5a
commit dc722adbb5
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 91 additions and 14 deletions

View File

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

View File

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