UniFi integration move to push messaging (#31086)

* Rewrite UniFi integration to use push messaging

* Add signalling for new clients/devices

* Update list of known wireless clients when we get events of them connecting

* Reconnection logic for websocket

* Fix failing tests

* Bump requirement to v12

* Add new tests

* Update homeassistant/components/unifi/controller.py

Co-Authored-By: Martin Hjelmare <marhje52@gmail.com>
This commit is contained in:
Robert Svensson 2020-01-31 20:23:25 +01:00 committed by GitHub
parent 06c8e53323
commit 958a867c11
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
14 changed files with 296 additions and 377 deletions

View File

@ -1,7 +1,7 @@
"""Support for devices connected to UniFi POE.""" """Support for devices connected to UniFi POE."""
import voluptuous as vol import voluptuous as vol
from homeassistant.const import CONF_HOST from homeassistant.const import CONF_HOST, EVENT_HOMEASSISTANT_STOP
from homeassistant.core import callback from homeassistant.core import callback
import homeassistant.helpers.config_validation as cv import homeassistant.helpers.config_validation as cv
from homeassistant.helpers.device_registry import CONNECTION_NETWORK_MAC from homeassistant.helpers.device_registry import CONNECTION_NETWORK_MAC
@ -96,6 +96,8 @@ async def async_setup_entry(hass, config_entry):
# sw_version=config.raw['swversion'], # sw_version=config.raw['swversion'],
) )
hass.bus.async_listen_once(EVENT_HOMEASSISTANT_STOP, controller.shutdown)
return True return True

View File

@ -5,9 +5,13 @@ import ssl
from aiohttp import CookieJar from aiohttp import CookieJar
import aiounifi import aiounifi
from aiounifi.controller import SIGNAL_CONNECTION_STATE
from aiounifi.events import WIRELESS_CLIENT_CONNECTED, WIRELESS_GUEST_CONNECTED
from aiounifi.websocket import STATE_DISCONNECTED, STATE_RUNNING
import async_timeout import async_timeout
from homeassistant.const import CONF_HOST from homeassistant.const import CONF_HOST
from homeassistant.core import callback
from homeassistant.exceptions import ConfigEntryNotReady from homeassistant.exceptions import ConfigEntryNotReady
from homeassistant.helpers import aiohttp_client from homeassistant.helpers import aiohttp_client
from homeassistant.helpers.dispatcher import async_dispatcher_send from homeassistant.helpers.dispatcher import async_dispatcher_send
@ -40,6 +44,7 @@ from .const import (
) )
from .errors import AuthenticationRequired, CannotConnect from .errors import AuthenticationRequired, CannotConnect
RETRY_TIMER = 15
SUPPORTED_PLATFORMS = ["device_tracker", "sensor", "switch"] SUPPORTED_PLATFORMS = ["device_tracker", "sensor", "switch"]
@ -59,6 +64,11 @@ class UniFiController:
self._site_name = None self._site_name = None
self._site_role = None self._site_role = None
@property
def controller_id(self):
"""Return the controller ID."""
return CONTROLLER_ID.format(host=self.host, site=self.site)
@property @property
def host(self): def host(self):
"""Return the host of this controller.""" """Return the host of this controller."""
@ -130,15 +140,47 @@ class UniFiController:
return client.mac return client.mac
return None return None
@callback
def async_unifi_signalling_callback(self, signal, data):
"""Handle messages back from UniFi library."""
if signal == SIGNAL_CONNECTION_STATE:
if data == STATE_DISCONNECTED and self.available:
LOGGER.error("Lost connection to UniFi")
if (data == STATE_RUNNING and not self.available) or (
data == STATE_DISCONNECTED and self.available
):
self.available = data == STATE_RUNNING
async_dispatcher_send(self.hass, self.signal_reachable)
if not self.available:
self.hass.loop.call_later(RETRY_TIMER, self.reconnect)
elif signal == "new_data" and data:
if "event" in data:
if data["event"].event in (
WIRELESS_CLIENT_CONNECTED,
WIRELESS_GUEST_CONNECTED,
):
self.update_wireless_clients()
else:
async_dispatcher_send(self.hass, self.signal_update)
@property
def signal_reachable(self) -> str:
"""Integration specific event to signal a change in connection status."""
return f"unifi-reachable-{self.controller_id}"
@property @property
def signal_update(self): def signal_update(self):
"""Event specific per UniFi entry to signal new data.""" """Event specific per UniFi entry to signal new data."""
return f"unifi-update-{CONTROLLER_ID.format(host=self.host, site=self.site)}" return f"unifi-update-{self.controller_id}"
@property @property
def signal_options_update(self): def signal_options_update(self):
"""Event specific per UniFi entry to signal new options.""" """Event specific per UniFi entry to signal new options."""
return f"unifi-options-{CONTROLLER_ID.format(host=self.host, site=self.site)}" return f"unifi-options-{self.controller_id}"
def update_wireless_clients(self): def update_wireless_clients(self):
"""Update set of known to be wireless clients.""" """Update set of known to be wireless clients."""
@ -156,59 +198,13 @@ class UniFiController:
unifi_wireless_clients = self.hass.data[UNIFI_WIRELESS_CLIENTS] unifi_wireless_clients = self.hass.data[UNIFI_WIRELESS_CLIENTS]
unifi_wireless_clients.update_data(self.wireless_clients, self.config_entry) unifi_wireless_clients.update_data(self.wireless_clients, self.config_entry)
async def request_update(self):
"""Request an update."""
if self.progress is not None:
return await self.progress
self.progress = self.hass.async_create_task(self.async_update())
await self.progress
self.progress = None
async def async_update(self):
"""Update UniFi controller information."""
failed = False
try:
with async_timeout.timeout(10):
await self.api.clients.update()
await self.api.devices.update()
if self.option_block_clients:
await self.api.clients_all.update()
except aiounifi.LoginRequired:
try:
with async_timeout.timeout(5):
await self.api.login()
except (asyncio.TimeoutError, aiounifi.AiounifiException):
failed = True
if self.available:
LOGGER.error("Unable to reach controller %s", self.host)
self.available = False
except (asyncio.TimeoutError, aiounifi.AiounifiException):
failed = True
if self.available:
LOGGER.error("Unable to reach controller %s", self.host)
self.available = False
if not failed and not self.available:
LOGGER.info("Reconnected to controller %s", self.host)
self.available = True
self.update_wireless_clients()
async_dispatcher_send(self.hass, self.signal_update)
async def async_setup(self): async def async_setup(self):
"""Set up a UniFi controller.""" """Set up a UniFi controller."""
hass = self.hass
try: try:
self.api = await get_controller( self.api = await get_controller(
self.hass, **self.config_entry.data[CONF_CONTROLLER] self.hass,
**self.config_entry.data[CONF_CONTROLLER],
async_callback=self.async_unifi_signalling_callback,
) )
await self.api.initialize() await self.api.initialize()
@ -227,21 +223,23 @@ class UniFiController:
LOGGER.error("Unknown error connecting with UniFi controller: %s", err) LOGGER.error("Unknown error connecting with UniFi controller: %s", err)
return False return False
wireless_clients = hass.data[UNIFI_WIRELESS_CLIENTS] wireless_clients = self.hass.data[UNIFI_WIRELESS_CLIENTS]
self.wireless_clients = wireless_clients.get_data(self.config_entry) self.wireless_clients = wireless_clients.get_data(self.config_entry)
self.update_wireless_clients() self.update_wireless_clients()
self.import_configuration() self.import_configuration()
self.config_entry.add_update_listener(self.async_options_updated)
for platform in SUPPORTED_PLATFORMS: for platform in SUPPORTED_PLATFORMS:
hass.async_create_task( self.hass.async_create_task(
hass.config_entries.async_forward_entry_setup( self.hass.config_entries.async_forward_entry_setup(
self.config_entry, platform self.config_entry, platform
) )
) )
self.api.start_websocket()
self.config_entry.add_update_listener(self.async_options_updated)
return True return True
@staticmethod @staticmethod
@ -296,12 +294,38 @@ class UniFiController:
self.config_entry, options=options self.config_entry, options=options
) )
@callback
def reconnect(self) -> None:
"""Prepare to reconnect UniFi session."""
LOGGER.debug("Reconnecting to UniFi in %i", RETRY_TIMER)
self.hass.loop.create_task(self.async_reconnect())
async def async_reconnect(self) -> None:
"""Try to reconnect UniFi session."""
try:
with async_timeout.timeout(5):
await self.api.login()
self.api.start_websocket()
except (asyncio.TimeoutError, aiounifi.AiounifiException):
self.hass.loop.call_later(RETRY_TIMER, self.reconnect)
@callback
def shutdown(self, event) -> None:
"""Wrap the call to unifi.close.
Used as an argument to EventBus.async_listen_once.
"""
self.api.stop_websocket()
async def async_reset(self): async def async_reset(self):
"""Reset this controller to default state. """Reset this controller to default state.
Will cancel any scheduled setup retry and will unload Will cancel any scheduled setup retry and will unload
the config entry. the config entry.
""" """
self.api.stop_websocket()
for platform in SUPPORTED_PLATFORMS: for platform in SUPPORTED_PLATFORMS:
await self.hass.config_entries.async_forward_entry_unload( await self.hass.config_entries.async_forward_entry_unload(
self.config_entry, platform self.config_entry, platform
@ -314,7 +338,9 @@ class UniFiController:
return True return True
async def get_controller(hass, host, username, password, port, site, verify_ssl): async def get_controller(
hass, host, username, password, port, site, verify_ssl, async_callback=None
):
"""Create a controller object and verify authentication.""" """Create a controller object and verify authentication."""
sslcontext = None sslcontext = None
@ -335,6 +361,7 @@ async def get_controller(hass, host, username, password, port, site, verify_ssl)
site=site, site=site,
websession=session, websession=session,
sslcontext=sslcontext, sslcontext=sslcontext,
callback=async_callback,
) )
try: try:

View File

@ -1,6 +1,5 @@
"""Track devices using UniFi controllers.""" """Track devices using UniFi controllers."""
import logging import logging
from pprint import pformat
from homeassistant.components.device_tracker import DOMAIN as DEVICE_TRACKER_DOMAIN from homeassistant.components.device_tracker import DOMAIN as DEVICE_TRACKER_DOMAIN
from homeassistant.components.device_tracker.config_entry import ScannerEntity from homeassistant.components.device_tracker.config_entry import ScannerEntity
@ -14,6 +13,7 @@ from homeassistant.helpers.entity_registry import DISABLED_CONFIG_ENTRY
import homeassistant.util.dt as dt_util import homeassistant.util.dt as dt_util
from .const import ATTR_MANUFACTURER from .const import ATTR_MANUFACTURER
from .unifi_client import UniFiClient
LOGGER = logging.getLogger(__name__) LOGGER = logging.getLogger(__name__)
@ -65,7 +65,7 @@ async def async_setup_entry(hass, config_entry, async_add_entities):
@callback @callback
def update_controller(): def update_controller():
"""Update the values of the controller.""" """Update the values of the controller."""
update_items(controller, async_add_entities, tracked) add_entities(controller, async_add_entities, tracked)
controller.listeners.append( controller.listeners.append(
async_dispatcher_connect(hass, controller.signal_update, update_controller) async_dispatcher_connect(hass, controller.signal_update, update_controller)
@ -97,8 +97,8 @@ async def async_setup_entry(hass, config_entry, async_add_entities):
@callback @callback
def update_items(controller, async_add_entities, tracked): def add_entities(controller, async_add_entities, tracked):
"""Update tracked device state from the controller.""" """Add new tracker entities from the controller."""
new_tracked = [] new_tracked = []
for items, tracker_class in ( for items, tracker_class in (
@ -109,8 +109,6 @@ def update_items(controller, async_add_entities, tracked):
for item_id in items: for item_id in items:
if item_id in tracked: if item_id in tracked:
if tracked[item_id].enabled:
tracked[item_id].async_schedule_update_ha_state()
continue continue
tracked[item_id] = tracker_class(items[item_id], controller) tracked[item_id] = tracker_class(items[item_id], controller)
@ -120,16 +118,14 @@ def update_items(controller, async_add_entities, tracked):
async_add_entities(new_tracked) async_add_entities(new_tracked)
class UniFiClientTracker(ScannerEntity): class UniFiClientTracker(UniFiClient, ScannerEntity):
"""Representation of a network client.""" """Representation of a network client."""
def __init__(self, client, controller): def __init__(self, client, controller):
"""Set up tracked client.""" """Set up tracked client."""
self.client = client super().__init__(client, controller)
self.controller = controller
self.is_wired = self.client.mac not in controller.wireless_clients
self.wired_bug = None
self.wired_bug = None
if self.is_wired != self.client.is_wired: if self.is_wired != self.client.is_wired:
self.wired_bug = dt_util.utcnow() - self.controller.option_detection_time self.wired_bug = dt_util.utcnow() - self.controller.option_detection_time
@ -151,26 +147,6 @@ class UniFiClientTracker(ScannerEntity):
return True return True
async def async_added_to_hass(self):
"""Client entity created."""
LOGGER.debug("New UniFi client tracker %s (%s)", self.name, self.client.mac)
async def async_update(self):
"""Synchronize state with controller.
Make sure to update self.is_wired if client is wireless, there is an issue when clients go offline that they get marked as wired.
"""
await self.controller.request_update()
if self.is_wired and self.client.mac in self.controller.wireless_clients:
self.is_wired = False
LOGGER.debug(
"Updating UniFi tracked client %s\n%s",
self.entity_id,
pformat(self.client.raw),
)
@property @property
def is_connected(self): def is_connected(self):
"""Return true if the client is connected to the network. """Return true if the client is connected to the network.
@ -198,26 +174,11 @@ class UniFiClientTracker(ScannerEntity):
"""Return the source type of the client.""" """Return the source type of the client."""
return SOURCE_TYPE_ROUTER return SOURCE_TYPE_ROUTER
@property
def name(self) -> str:
"""Return the name of the client."""
return self.client.name or self.client.hostname
@property @property
def unique_id(self) -> str: def unique_id(self) -> str:
"""Return a unique identifier for this client.""" """Return a unique identifier for this client."""
return f"{self.client.mac}-{self.controller.site}" return f"{self.client.mac}-{self.controller.site}"
@property
def available(self) -> bool:
"""Return if controller is available."""
return self.controller.available
@property
def device_info(self):
"""Return a client description for device registry."""
return {"connections": {(CONNECTION_NETWORK_MAC, self.client.mac)}}
@property @property
def device_state_attributes(self): def device_state_attributes(self):
"""Return the client state attributes.""" """Return the client state attributes."""
@ -239,6 +200,7 @@ class UniFiDeviceTracker(ScannerEntity):
"""Set up tracked device.""" """Set up tracked device."""
self.device = device self.device = device
self.controller = controller self.controller = controller
self.listeners = []
@property @property
def entity_registry_enabled_default(self): def entity_registry_enabled_default(self):
@ -250,17 +212,26 @@ class UniFiDeviceTracker(ScannerEntity):
async def async_added_to_hass(self): async def async_added_to_hass(self):
"""Subscribe to device events.""" """Subscribe to device events."""
LOGGER.debug("New UniFi device tracker %s (%s)", self.name, self.device.mac) LOGGER.debug("New UniFi device tracker %s (%s)", self.name, self.device.mac)
self.device.register_callback(self.async_update_callback)
async def async_update(self): self.listeners.append(
"""Synchronize state with controller.""" async_dispatcher_connect(
await self.controller.request_update() self.hass, self.controller.signal_reachable, self.async_update_callback
)
LOGGER.debug(
"Updating UniFi tracked device %s\n%s",
self.entity_id,
pformat(self.device.raw),
) )
async def async_will_remove_from_hass(self) -> None:
"""Disconnect device object when removed."""
self.device.remove_callback(self.async_update_callback)
for unsub_dispatcher in self.listeners:
unsub_dispatcher()
@callback
def async_update_callback(self):
"""Update the sensor's state."""
LOGGER.debug("Updating UniFi tracked device %s", self.entity_id)
self.async_schedule_update_ha_state()
@property @property
def is_connected(self): def is_connected(self):
"""Return true if the device is connected to the network.""" """Return true if the device is connected to the network."""
@ -325,3 +296,8 @@ class UniFiDeviceTracker(ScannerEntity):
attributes["upgradable"] = self.device.upgradable attributes["upgradable"] = self.device.upgradable
return attributes return attributes
@property
def should_poll(self):
"""No polling needed."""
return False

View File

@ -3,8 +3,12 @@
"name": "Ubiquiti UniFi", "name": "Ubiquiti UniFi",
"config_flow": true, "config_flow": true,
"documentation": "https://www.home-assistant.io/integrations/unifi", "documentation": "https://www.home-assistant.io/integrations/unifi",
"requirements": ["aiounifi==11"], "requirements": [
"aiounifi==12"
],
"dependencies": [], "dependencies": [],
"codeowners": ["@kane610"], "codeowners": [
"@kane610"
],
"quality_scale": "platinum" "quality_scale": "platinum"
} }

View File

@ -4,11 +4,11 @@ import logging
from homeassistant.components.unifi.config_flow import get_controller_from_config_entry from homeassistant.components.unifi.config_flow import get_controller_from_config_entry
from homeassistant.core import callback from homeassistant.core import callback
from homeassistant.helpers import entity_registry 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.dispatcher import async_dispatcher_connect
from homeassistant.helpers.entity import Entity
from homeassistant.helpers.entity_registry import DISABLED_CONFIG_ENTRY from homeassistant.helpers.entity_registry import DISABLED_CONFIG_ENTRY
from .unifi_client import UniFiClient
LOGGER = logging.getLogger(__name__) LOGGER = logging.getLogger(__name__)
ATTR_RECEIVING = "receiving" ATTR_RECEIVING = "receiving"
@ -29,7 +29,7 @@ async def async_setup_entry(hass, config_entry, async_add_entities):
@callback @callback
def update_controller(): def update_controller():
"""Update the values of the controller.""" """Update the values of the controller."""
update_items(controller, async_add_entities, sensors) add_entities(controller, async_add_entities, sensors)
controller.listeners.append( controller.listeners.append(
async_dispatcher_connect(hass, controller.signal_update, update_controller) async_dispatcher_connect(hass, controller.signal_update, update_controller)
@ -61,8 +61,8 @@ async def async_setup_entry(hass, config_entry, async_add_entities):
@callback @callback
def update_items(controller, async_add_entities, sensors): def add_entities(controller, async_add_entities, sensors):
"""Update sensors from the controller.""" """Add new sensor entities from the controller."""
new_sensors = [] new_sensors = []
for client_id in controller.api.clients: for client_id in controller.api.clients:
@ -73,9 +73,6 @@ def update_items(controller, async_add_entities, sensors):
item_id = f"{direction}-{client_id}" item_id = f"{direction}-{client_id}"
if item_id in sensors: if item_id in sensors:
sensor = sensors[item_id]
if sensor.enabled:
sensor.async_schedule_update_ha_state()
continue continue
sensors[item_id] = sensor_class( sensors[item_id] = sensor_class(
@ -87,14 +84,8 @@ def update_items(controller, async_add_entities, sensors):
async_add_entities(new_sensors) async_add_entities(new_sensors)
class UniFiBandwidthSensor(Entity): class UniFiRxBandwidthSensor(UniFiClient):
"""UniFi Bandwidth sensor base class.""" """Receiving bandwidth sensor."""
def __init__(self, client, controller):
"""Set up client."""
self.client = client
self.controller = controller
self.is_wired = self.client.mac not in controller.wireless_clients
@property @property
def entity_registry_enabled_default(self): def entity_registry_enabled_default(self):
@ -103,37 +94,6 @@ class UniFiBandwidthSensor(Entity):
return True return True
return False return False
async def async_added_to_hass(self):
"""Client entity created."""
LOGGER.debug("New UniFi bandwidth sensor %s (%s)", self.name, self.client.mac)
async def async_update(self):
"""Synchronize state with controller.
Make sure to update self.is_wired if client is wireless, there is an issue when clients go offline that they get marked as wired.
"""
LOGGER.debug(
"Updating UniFi bandwidth sensor %s (%s)", self.entity_id, self.client.mac
)
await self.controller.request_update()
if self.is_wired and self.client.mac in self.controller.wireless_clients:
self.is_wired = False
@property
def available(self) -> bool:
"""Return if controller is available."""
return self.controller.available
@property
def device_info(self):
"""Return a device description for device registry."""
return {"connections": {(CONNECTION_NETWORK_MAC, self.client.mac)}}
class UniFiRxBandwidthSensor(UniFiBandwidthSensor):
"""Receiving bandwidth sensor."""
@property @property
def state(self): def state(self):
"""Return the state of the sensor.""" """Return the state of the sensor."""
@ -153,7 +113,7 @@ class UniFiRxBandwidthSensor(UniFiBandwidthSensor):
return f"rx-{self.client.mac}" return f"rx-{self.client.mac}"
class UniFiTxBandwidthSensor(UniFiBandwidthSensor): class UniFiTxBandwidthSensor(UniFiRxBandwidthSensor):
"""Transmitting bandwidth sensor.""" """Transmitting bandwidth sensor."""
@property @property

View File

@ -1,15 +1,15 @@
"""Support for devices connected to UniFi POE.""" """Support for devices connected to UniFi POE."""
import logging import logging
from pprint import pformat
from homeassistant.components.switch import SwitchDevice from homeassistant.components.switch import SwitchDevice
from homeassistant.components.unifi.config_flow import get_controller_from_config_entry from homeassistant.components.unifi.config_flow import get_controller_from_config_entry
from homeassistant.core import callback from homeassistant.core import callback
from homeassistant.helpers import entity_registry 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.dispatcher import async_dispatcher_connect
from homeassistant.helpers.restore_state import RestoreEntity from homeassistant.helpers.restore_state import RestoreEntity
from .unifi_client import UniFiClient
LOGGER = logging.getLogger(__name__) LOGGER = logging.getLogger(__name__)
@ -20,7 +20,7 @@ async def async_setup_platform(hass, config, async_add_entities, discovery_info=
async def async_setup_entry(hass, config_entry, async_add_entities): async def async_setup_entry(hass, config_entry, async_add_entities):
"""Set up switches for UniFi component. """Set up switches for UniFi component.
Switches are controlling network switch ports with Poe. Switches are controlling network access and switch ports with POE.
""" """
controller = get_controller_from_config_entry(hass, config_entry) controller = get_controller_from_config_entry(hass, config_entry)
@ -55,7 +55,7 @@ async def async_setup_entry(hass, config_entry, async_add_entities):
@callback @callback
def update_controller(): def update_controller():
"""Update the values of the controller.""" """Update the values of the controller."""
update_items(controller, async_add_entities, switches, switches_off) add_entities(controller, async_add_entities, switches, switches_off)
controller.listeners.append( controller.listeners.append(
async_dispatcher_connect(hass, controller.signal_update, update_controller) async_dispatcher_connect(hass, controller.signal_update, update_controller)
@ -66,8 +66,8 @@ async def async_setup_entry(hass, config_entry, async_add_entities):
@callback @callback
def update_items(controller, async_add_entities, switches, switches_off): def add_entities(controller, async_add_entities, switches, switches_off):
"""Update POE port state from the controller.""" """Add new switch entities from the controller."""
new_switches = [] new_switches = []
devices = controller.api.devices devices = controller.api.devices
@ -77,13 +77,6 @@ def update_items(controller, async_add_entities, switches, switches_off):
block_client_id = f"block-{client_id}" block_client_id = f"block-{client_id}"
if block_client_id in switches: if block_client_id in switches:
if switches[block_client_id].enabled:
LOGGER.debug(
"Updating UniFi block switch %s (%s)",
switches[block_client_id].entity_id,
switches[block_client_id].client.mac,
)
switches[block_client_id].async_schedule_update_ha_state()
continue continue
if client_id not in controller.api.clients_all: if client_id not in controller.api.clients_all:
@ -99,13 +92,6 @@ def update_items(controller, async_add_entities, switches, switches_off):
poe_client_id = f"poe-{client_id}" poe_client_id = f"poe-{client_id}"
if poe_client_id in switches: if poe_client_id in switches:
if switches[poe_client_id].enabled:
LOGGER.debug(
"Updating UniFi POE switch %s (%s)",
switches[poe_client_id].entity_id,
switches[poe_client_id].client.mac,
)
switches[poe_client_id].async_schedule_update_ha_state()
continue continue
client = controller.api.clients[client_id] client = controller.api.clients[client_id]
@ -148,42 +134,21 @@ def update_items(controller, async_add_entities, switches, switches_off):
async_add_entities(new_switches) async_add_entities(new_switches)
class UniFiClient:
"""Base class for UniFi switches."""
def __init__(self, client, controller):
"""Set up switch."""
self.client = client
self.controller = controller
async def async_update(self):
"""Synchronize state with controller."""
await self.controller.request_update()
@property
def name(self):
"""Return the name of the client."""
return self.client.name or self.client.hostname
@property
def device_info(self):
"""Return a device description for device registry."""
return {"connections": {(CONNECTION_NETWORK_MAC, self.client.mac)}}
class UniFiPOEClientSwitch(UniFiClient, SwitchDevice, RestoreEntity): class UniFiPOEClientSwitch(UniFiClient, SwitchDevice, RestoreEntity):
"""Representation of a client that uses POE.""" """Representation of a client that uses POE."""
def __init__(self, client, controller): def __init__(self, client, controller):
"""Set up POE switch.""" """Set up POE switch."""
super().__init__(client, controller) super().__init__(client, controller)
self.poe_mode = None self.poe_mode = None
if self.client.sw_port and self.port.poe_mode != "off": if self.client.sw_port and self.port.poe_mode != "off":
self.poe_mode = self.port.poe_mode self.poe_mode = self.port.poe_mode
async def async_added_to_hass(self): async def async_added_to_hass(self):
"""Call when entity about to be added to Home Assistant.""" """Call when entity about to be added to Home Assistant."""
LOGGER.debug("New UniFi POE switch %s (%s)", self.name, self.client.mac) await super().async_added_to_hass()
state = await self.async_get_last_state() state = await self.async_get_last_state()
if state is None: if state is None:
@ -198,16 +163,6 @@ class UniFiPOEClientSwitch(UniFiClient, SwitchDevice, RestoreEntity):
if not self.client.sw_port: if not self.client.sw_port:
self.client.raw["sw_port"] = state.attributes["port"] self.client.raw["sw_port"] = state.attributes["port"]
async def async_update(self):
"""Log client information after update."""
await super().async_update()
LOGGER.debug(
"Updating UniFi POE controlled client %s\n%s",
self.entity_id,
pformat(self.client.raw),
)
@property @property
def unique_id(self): def unique_id(self):
"""Return a unique identifier for this switch.""" """Return a unique identifier for this switch."""
@ -267,10 +222,6 @@ class UniFiPOEClientSwitch(UniFiClient, SwitchDevice, RestoreEntity):
class UniFiBlockClientSwitch(UniFiClient, SwitchDevice): class UniFiBlockClientSwitch(UniFiClient, SwitchDevice):
"""Representation of a blockable client.""" """Representation of a blockable client."""
async def async_added_to_hass(self):
"""Call when entity about to be added to Home Assistant."""
LOGGER.debug("New UniFi Block switch %s (%s)", self.name, self.client.mac)
@property @property
def unique_id(self): def unique_id(self):
"""Return a unique identifier for this switch.""" """Return a unique identifier for this switch."""
@ -281,11 +232,6 @@ class UniFiBlockClientSwitch(UniFiClient, SwitchDevice):
"""Return true if client is allowed to connect.""" """Return true if client is allowed to connect."""
return not self.client.blocked return not self.client.blocked
@property
def available(self):
"""Return if controller is available."""
return self.controller.available
async def async_turn_on(self, **kwargs): async def async_turn_on(self, **kwargs):
"""Turn on connectivity for client.""" """Turn on connectivity for client."""
await self.controller.api.clients.async_unblock(self.client.mac) await self.controller.api.clients.async_unblock(self.client.mac)

View File

@ -0,0 +1,65 @@
"""Base class for UniFi clients."""
import logging
from homeassistant.core import callback
from homeassistant.helpers.device_registry import CONNECTION_NETWORK_MAC
from homeassistant.helpers.dispatcher import async_dispatcher_connect
from homeassistant.helpers.entity import Entity
LOGGER = logging.getLogger(__name__)
class UniFiClient(Entity):
"""Base class for UniFi clients."""
def __init__(self, client, controller) -> None:
"""Set up client."""
self.client = client
self.controller = controller
self.listeners = []
self.is_wired = self.client.mac not in controller.wireless_clients
async def async_added_to_hass(self) -> None:
"""Client entity created."""
LOGGER.debug("New UniFi client %s (%s)", self.name, self.client.mac)
self.client.register_callback(self.async_update_callback)
self.listeners.append(
async_dispatcher_connect(
self.hass, self.controller.signal_reachable, self.async_update_callback
)
)
async def async_will_remove_from_hass(self) -> None:
"""Disconnect client object when removed."""
self.client.remove_callback(self.async_update_callback)
for unsub_dispatcher in self.listeners:
unsub_dispatcher()
@callback
def async_update_callback(self) -> None:
"""Update the clients state."""
if self.is_wired and self.client.mac in self.controller.wireless_clients:
self.is_wired = False
LOGGER.debug("Updating client %s %s", self.entity_id, self.client.mac)
self.async_schedule_update_ha_state()
@property
def name(self) -> str:
"""Return the name of the client."""
return self.client.name or self.client.hostname
@property
def available(self) -> bool:
"""Return if controller is available."""
return self.controller.available
@property
def device_info(self) -> dict:
"""Return a client description for device registry."""
return {"connections": {(CONNECTION_NETWORK_MAC, self.client.mac)}}
@property
def should_poll(self) -> bool:
"""No polling needed."""
return False

View File

@ -196,7 +196,7 @@ aiopylgtv==0.3.2
aioswitcher==2019.4.26 aioswitcher==2019.4.26
# homeassistant.components.unifi # homeassistant.components.unifi
aiounifi==11 aiounifi==12
# homeassistant.components.wwlln # homeassistant.components.wwlln
aiowwlln==2.0.2 aiowwlln==2.0.2

View File

@ -75,7 +75,7 @@ aiopylgtv==0.3.2
aioswitcher==2019.4.26 aioswitcher==2019.4.26
# homeassistant.components.unifi # homeassistant.components.unifi
aiounifi==11 aiounifi==12
# homeassistant.components.wwlln # homeassistant.components.wwlln
aiowwlln==2.0.2 aiowwlln==2.0.2

View File

@ -111,9 +111,12 @@ async def setup_unifi_integration(
return mock_client_all_responses.popleft() return mock_client_all_responses.popleft()
return {} return {}
# "aiounifi.Controller.start_websocket", return_value=True
with patch("aiounifi.Controller.login", return_value=True), patch( with patch("aiounifi.Controller.login", return_value=True), patch(
"aiounifi.Controller.sites", return_value=sites "aiounifi.Controller.sites", return_value=sites
), patch("aiounifi.Controller.request", new=mock_request): ), patch("aiounifi.Controller.request", new=mock_request), patch.object(
aiounifi.websocket.WSClient, "start", return_value=True
):
await hass.config_entries.async_setup(config_entry.entry_id) await hass.config_entries.async_setup(config_entry.entry_id)
await hass.async_block_till_done() await hass.async_block_till_done()
@ -233,47 +236,28 @@ async def test_reset_after_successful_setup(hass):
assert len(controller.listeners) == 0 assert len(controller.listeners) == 0
async def test_failed_update_failed_login(hass): async def test_wireless_client_event_calls_update_wireless_devices(hass):
"""Running update can handle a failed login.""" """Call update_wireless_devices method when receiving wireless client event."""
controller = await setup_unifi_integration(hass) controller = await setup_unifi_integration(hass)
with patch.object( with patch(
controller.api.clients, "update", side_effect=aiounifi.LoginRequired "homeassistant.components.unifi.controller.UniFiController.update_wireless_clients",
), patch.object(controller.api, "login", side_effect=aiounifi.AiounifiException): return_value=None,
await controller.async_update() ) as wireless_clients_mock:
await hass.async_block_till_done() controller.api.websocket._data = {
"meta": {"rc": "ok", "message": "events"},
"data": [
{
"datetime": "2020-01-20T19:37:04Z",
"key": aiounifi.events.WIRELESS_CLIENT_CONNECTED,
"msg": "User[11:22:33:44:55:66] has connected to WLAN",
"time": 1579549024893,
}
],
}
controller.api.session_handler("data")
assert controller.available is False assert wireless_clients_mock.assert_called_once
async def test_failed_update_successful_login(hass):
"""Running update can login when requested."""
controller = await setup_unifi_integration(hass)
with patch.object(
controller.api.clients, "update", side_effect=aiounifi.LoginRequired
), patch.object(controller.api, "login", return_value=Mock(True)):
await controller.async_update()
await hass.async_block_till_done()
assert controller.available is True
async def test_failed_update(hass):
"""Running update can login when requested."""
controller = await setup_unifi_integration(hass)
with patch.object(
controller.api.clients, "update", side_effect=aiounifi.AiounifiException
):
await controller.async_update()
await hass.async_block_till_done()
assert controller.available is False
await controller.async_update()
await hass.async_block_till_done()
assert controller.available is True
async def test_get_controller(hass): async def test_get_controller(hass):

View File

@ -2,6 +2,8 @@
from copy import copy from copy import copy
from datetime import timedelta from datetime import timedelta
from aiounifi.controller import SIGNAL_CONNECTION_STATE
from aiounifi.websocket import STATE_DISCONNECTED, STATE_RUNNING
from asynctest import patch from asynctest import patch
from homeassistant import config_entries from homeassistant import config_entries
@ -136,11 +138,12 @@ async def test_tracked_devices(hass):
client_1_copy = copy(CLIENT_1) client_1_copy = copy(CLIENT_1)
client_1_copy["last_seen"] = dt_util.as_timestamp(dt_util.utcnow()) client_1_copy["last_seen"] = dt_util.as_timestamp(dt_util.utcnow())
event = {"meta": {"message": "sta:sync"}, "data": [client_1_copy]}
controller.api.message_handler(event)
device_1_copy = copy(DEVICE_1) device_1_copy = copy(DEVICE_1)
device_1_copy["last_seen"] = dt_util.as_timestamp(dt_util.utcnow()) device_1_copy["last_seen"] = dt_util.as_timestamp(dt_util.utcnow())
controller.mock_client_responses.append([client_1_copy]) event = {"meta": {"message": "device:sync"}, "data": [device_1_copy]}
controller.mock_device_responses.append([device_1_copy]) controller.api.message_handler(event)
await controller.async_update()
await hass.async_block_till_done() await hass.async_block_till_done()
client_1 = hass.states.get("device_tracker.client_1") client_1 = hass.states.get("device_tracker.client_1")
@ -149,16 +152,39 @@ async def test_tracked_devices(hass):
device_1 = hass.states.get("device_tracker.device_1") device_1 = hass.states.get("device_tracker.device_1")
assert device_1.state == "home" assert device_1.state == "home"
# Controller unavailable
controller.async_unifi_signalling_callback(
SIGNAL_CONNECTION_STATE, STATE_DISCONNECTED
)
await hass.async_block_till_done()
client_1 = hass.states.get("device_tracker.client_1")
assert client_1.state == STATE_UNAVAILABLE
device_1 = hass.states.get("device_tracker.device_1")
assert device_1.state == STATE_UNAVAILABLE
# Controller available
controller.async_unifi_signalling_callback(SIGNAL_CONNECTION_STATE, STATE_RUNNING)
await hass.async_block_till_done()
client_1 = hass.states.get("device_tracker.client_1")
assert client_1.state == "home"
device_1 = hass.states.get("device_tracker.device_1")
assert device_1.state == "home"
# Disabled device is unavailable
device_1_copy = copy(DEVICE_1) device_1_copy = copy(DEVICE_1)
device_1_copy["disabled"] = True device_1_copy["disabled"] = True
controller.mock_client_responses.append({}) event = {"meta": {"message": "device:sync"}, "data": [device_1_copy]}
controller.mock_device_responses.append([device_1_copy]) controller.api.message_handler(event)
await controller.async_update()
await hass.async_block_till_done() await hass.async_block_till_done()
device_1 = hass.states.get("device_tracker.device_1") device_1 = hass.states.get("device_tracker.device_1")
assert device_1.state == STATE_UNAVAILABLE assert device_1.state == STATE_UNAVAILABLE
# Don't track wired clients nor devices
controller.config_entry.add_update_listener(controller.async_options_updated) controller.config_entry.add_update_listener(controller.async_options_updated)
hass.config_entries.async_update_entry( hass.config_entries.async_update_entry(
controller.config_entry, controller.config_entry,
@ -194,9 +220,8 @@ async def test_wireless_client_go_wired_issue(hass):
client_1_client["is_wired"] = True client_1_client["is_wired"] = True
client_1_client["last_seen"] = dt_util.as_timestamp(dt_util.utcnow()) client_1_client["last_seen"] = dt_util.as_timestamp(dt_util.utcnow())
controller.mock_client_responses.append([client_1_client]) event = {"meta": {"message": "sta:sync"}, "data": [client_1_client]}
controller.mock_device_responses.append({}) controller.api.message_handler(event)
await controller.async_update()
await hass.async_block_till_done() await hass.async_block_till_done()
client_1 = hass.states.get("device_tracker.client_1") client_1 = hass.states.get("device_tracker.client_1")
@ -207,9 +232,8 @@ async def test_wireless_client_go_wired_issue(hass):
"utcnow", "utcnow",
return_value=(dt_util.utcnow() + timedelta(minutes=5)), return_value=(dt_util.utcnow() + timedelta(minutes=5)),
): ):
controller.mock_client_responses.append([client_1_client]) event = {"meta": {"message": "sta:sync"}, "data": [client_1_client]}
controller.mock_device_responses.append({}) controller.api.message_handler(event)
await controller.async_update()
await hass.async_block_till_done() await hass.async_block_till_done()
client_1 = hass.states.get("device_tracker.client_1") client_1 = hass.states.get("device_tracker.client_1")
@ -217,9 +241,8 @@ async def test_wireless_client_go_wired_issue(hass):
client_1_client["is_wired"] = False client_1_client["is_wired"] = False
client_1_client["last_seen"] = dt_util.as_timestamp(dt_util.utcnow()) client_1_client["last_seen"] = dt_util.as_timestamp(dt_util.utcnow())
controller.mock_client_responses.append([client_1_client]) event = {"meta": {"message": "sta:sync"}, "data": [client_1_client]}
controller.mock_device_responses.append({}) controller.api.message_handler(event)
await controller.async_update()
await hass.async_block_till_done() await hass.async_block_till_done()
client_1 = hass.states.get("device_tracker.client_1") client_1 = hass.states.get("device_tracker.client_1")

View File

@ -4,6 +4,8 @@ from unittest.mock import Mock, patch
from homeassistant.components import unifi from homeassistant.components import unifi
from homeassistant.setup import async_setup_component from homeassistant.setup import async_setup_component
from .test_controller import setup_unifi_integration
from tests.common import MockConfigEntry, mock_coro from tests.common import MockConfigEntry, mock_coro
@ -42,67 +44,15 @@ async def test_setup_with_config(hass):
async def test_successful_config_entry(hass): async def test_successful_config_entry(hass):
"""Test that configured options for a host are loaded via config entry.""" """Test that configured options for a host are loaded via config entry."""
entry = MockConfigEntry( await setup_unifi_integration(hass)
domain=unifi.DOMAIN, assert hass.data[unifi.DOMAIN]
data={
"controller": {
"host": "0.0.0.0",
"username": "user",
"password": "pass",
"port": 80,
"site": "default",
"verify_ssl": True,
},
"poe_control": True,
},
)
entry.add_to_hass(hass)
mock_registry = Mock()
with patch.object(unifi, "UniFiController") as mock_controller, patch(
"homeassistant.helpers.device_registry.async_get_registry",
return_value=mock_coro(mock_registry),
):
mock_controller.return_value.async_setup.return_value = mock_coro(True)
mock_controller.return_value.mac = "00:11:22:33:44:55"
assert await unifi.async_setup_entry(hass, entry) is True
assert len(mock_controller.mock_calls) == 2
p_hass, p_entry = mock_controller.mock_calls[0][1]
assert p_hass is hass
assert p_entry is entry
assert len(mock_registry.mock_calls) == 1
assert mock_registry.mock_calls[0][2] == {
"config_entry_id": entry.entry_id,
"connections": {("mac", "00:11:22:33:44:55")},
"manufacturer": unifi.ATTR_MANUFACTURER,
"model": "UniFi Controller",
"name": "UniFi Controller",
}
async def test_controller_fail_setup(hass): async def test_controller_fail_setup(hass):
"""Test that a failed setup still stores controller.""" """Test that a failed setup still stores controller."""
entry = MockConfigEntry(
domain=unifi.DOMAIN,
data={
"controller": {
"host": "0.0.0.0",
"username": "user",
"password": "pass",
"port": 80,
"site": "default",
"verify_ssl": True,
},
"poe_control": True,
},
)
entry.add_to_hass(hass)
with patch.object(unifi, "UniFiController") as mock_cntrlr: with patch.object(unifi, "UniFiController") as mock_cntrlr:
mock_cntrlr.return_value.async_setup.return_value = mock_coro(False) mock_cntrlr.return_value.async_setup.return_value = mock_coro(False)
assert await unifi.async_setup_entry(hass, entry) is False await setup_unifi_integration(hass)
assert hass.data[unifi.DOMAIN] == {} assert hass.data[unifi.DOMAIN] == {}
@ -140,33 +90,8 @@ async def test_controller_no_mac(hass):
async def test_unload_entry(hass): async def test_unload_entry(hass):
"""Test being able to unload an entry.""" """Test being able to unload an entry."""
entry = MockConfigEntry( controller = await setup_unifi_integration(hass)
domain=unifi.DOMAIN, assert hass.data[unifi.DOMAIN]
data={
"controller": {
"host": "0.0.0.0",
"username": "user",
"password": "pass",
"port": 80,
"site": "default",
"verify_ssl": True,
},
"poe_control": True,
},
)
entry.add_to_hass(hass)
with patch.object(unifi, "UniFiController") as mock_controller, patch( assert await unifi.async_unload_entry(hass, controller.config_entry)
"homeassistant.helpers.device_registry.async_get_registry", assert not hass.data[unifi.DOMAIN]
return_value=mock_coro(Mock()),
):
mock_controller.return_value.async_setup.return_value = mock_coro(True)
mock_controller.return_value.mac = "00:11:22:33:44:55"
assert await unifi.async_setup_entry(hass, entry) is True
assert len(mock_controller.return_value.mock_calls) == 1
mock_controller.return_value.async_reset.return_value = mock_coro(True)
assert await unifi.async_unload_entry(hass, entry)
assert len(mock_controller.return_value.async_reset.mock_calls) == 1
assert hass.data[unifi.DOMAIN] == {}

View File

@ -90,8 +90,8 @@ async def test_sensors(hass):
clients[1]["rx_bytes"] = 2345000000 clients[1]["rx_bytes"] = 2345000000
clients[1]["tx_bytes"] = 6789000000 clients[1]["tx_bytes"] = 6789000000
controller.mock_client_responses.append(clients) event = {"meta": {"message": "sta:sync"}, "data": clients}
await controller.async_update() controller.api.message_handler(event)
await hass.async_block_till_done() await hass.async_block_till_done()
wireless_client_rx = hass.states.get("sensor.wireless_client_name_rx") wireless_client_rx = hass.states.get("sensor.wireless_client_name_rx")

View File

@ -300,13 +300,17 @@ async def test_new_client_discovered_on_block_control(hass):
assert len(controller.mock_requests) == 3 assert len(controller.mock_requests) == 3
assert len(hass.states.async_all()) == 2 assert len(hass.states.async_all()) == 2
controller.mock_client_all_responses.append([BLOCKED]) controller.api.websocket._data = {
"meta": {"message": "sta:sync"},
"data": [BLOCKED],
}
controller.api.session_handler("data")
# Calling a service will trigger the updates to run # Calling a service will trigger the updates to run
await hass.services.async_call( await hass.services.async_call(
"switch", "turn_off", {"entity_id": "switch.block_client_1"}, blocking=True "switch", "turn_off", {"entity_id": "switch.block_client_1"}, blocking=True
) )
assert len(controller.mock_requests) == 7 assert len(controller.mock_requests) == 4
assert len(hass.states.async_all()) == 2 assert len(hass.states.async_all()) == 2
assert controller.mock_requests[3] == { assert controller.mock_requests[3] == {
"json": {"mac": "00:00:00:00:01:01", "cmd": "block-sta"}, "json": {"mac": "00:00:00:00:01:01", "cmd": "block-sta"},
@ -317,8 +321,8 @@ async def test_new_client_discovered_on_block_control(hass):
await hass.services.async_call( await hass.services.async_call(
"switch", "turn_on", {"entity_id": "switch.block_client_1"}, blocking=True "switch", "turn_on", {"entity_id": "switch.block_client_1"}, blocking=True
) )
assert len(controller.mock_requests) == 11 assert len(controller.mock_requests) == 5
assert controller.mock_requests[7] == { assert controller.mock_requests[4] == {
"json": {"mac": "00:00:00:00:01:01", "cmd": "unblock-sta"}, "json": {"mac": "00:00:00:00:01:01", "cmd": "unblock-sta"},
"method": "post", "method": "post",
"path": "s/{site}/cmd/stamgr/", "path": "s/{site}/cmd/stamgr/",
@ -340,14 +344,17 @@ async def test_new_client_discovered_on_poe_control(hass):
assert len(controller.mock_requests) == 3 assert len(controller.mock_requests) == 3
assert len(hass.states.async_all()) == 2 assert len(hass.states.async_all()) == 2
controller.mock_client_responses.append([CLIENT_1, CLIENT_2]) controller.api.websocket._data = {
controller.mock_device_responses.append([DEVICE_1]) "meta": {"message": "sta:sync"},
"data": [CLIENT_2],
}
controller.api.session_handler("data")
# Calling a service will trigger the updates to run # Calling a service will trigger the updates to run
await hass.services.async_call( await hass.services.async_call(
"switch", "turn_off", {"entity_id": "switch.poe_client_1"}, blocking=True "switch", "turn_off", {"entity_id": "switch.poe_client_1"}, blocking=True
) )
assert len(controller.mock_requests) == 6 assert len(controller.mock_requests) == 4
assert len(hass.states.async_all()) == 3 assert len(hass.states.async_all()) == 3
assert controller.mock_requests[3] == { assert controller.mock_requests[3] == {
"json": { "json": {
@ -360,7 +367,7 @@ async def test_new_client_discovered_on_poe_control(hass):
await hass.services.async_call( await hass.services.async_call(
"switch", "turn_on", {"entity_id": "switch.poe_client_1"}, blocking=True "switch", "turn_on", {"entity_id": "switch.poe_client_1"}, blocking=True
) )
assert len(controller.mock_requests) == 9 assert len(controller.mock_requests) == 5
assert controller.mock_requests[3] == { assert controller.mock_requests[3] == {
"json": { "json": {
"port_overrides": [ "port_overrides": [