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."""
import voluptuous as vol
from homeassistant.const import CONF_HOST
from homeassistant.const import CONF_HOST, EVENT_HOMEASSISTANT_STOP
from homeassistant.core import callback
import homeassistant.helpers.config_validation as cv
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'],
)
hass.bus.async_listen_once(EVENT_HOMEASSISTANT_STOP, controller.shutdown)
return True

View File

@ -5,9 +5,13 @@ import ssl
from aiohttp import CookieJar
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
from homeassistant.const import CONF_HOST
from homeassistant.core import callback
from homeassistant.exceptions import ConfigEntryNotReady
from homeassistant.helpers import aiohttp_client
from homeassistant.helpers.dispatcher import async_dispatcher_send
@ -40,6 +44,7 @@ from .const import (
)
from .errors import AuthenticationRequired, CannotConnect
RETRY_TIMER = 15
SUPPORTED_PLATFORMS = ["device_tracker", "sensor", "switch"]
@ -59,6 +64,11 @@ class UniFiController:
self._site_name = 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
def host(self):
"""Return the host of this controller."""
@ -130,15 +140,47 @@ class UniFiController:
return client.mac
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
def signal_update(self):
"""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
def signal_options_update(self):
"""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):
"""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.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):
"""Set up a UniFi controller."""
hass = self.hass
try:
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()
@ -227,21 +223,23 @@ class UniFiController:
LOGGER.error("Unknown error connecting with UniFi controller: %s", err)
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.update_wireless_clients()
self.import_configuration()
self.config_entry.add_update_listener(self.async_options_updated)
for platform in SUPPORTED_PLATFORMS:
hass.async_create_task(
hass.config_entries.async_forward_entry_setup(
self.hass.async_create_task(
self.hass.config_entries.async_forward_entry_setup(
self.config_entry, platform
)
)
self.api.start_websocket()
self.config_entry.add_update_listener(self.async_options_updated)
return True
@staticmethod
@ -296,12 +294,38 @@ class UniFiController:
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):
"""Reset this controller to default state.
Will cancel any scheduled setup retry and will unload
the config entry.
"""
self.api.stop_websocket()
for platform in SUPPORTED_PLATFORMS:
await self.hass.config_entries.async_forward_entry_unload(
self.config_entry, platform
@ -314,7 +338,9 @@ class UniFiController:
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."""
sslcontext = None
@ -335,6 +361,7 @@ async def get_controller(hass, host, username, password, port, site, verify_ssl)
site=site,
websession=session,
sslcontext=sslcontext,
callback=async_callback,
)
try:

View File

@ -1,6 +1,5 @@
"""Track devices using UniFi controllers."""
import logging
from pprint import pformat
from homeassistant.components.device_tracker import DOMAIN as DEVICE_TRACKER_DOMAIN
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
from .const import ATTR_MANUFACTURER
from .unifi_client import UniFiClient
LOGGER = logging.getLogger(__name__)
@ -65,7 +65,7 @@ async def async_setup_entry(hass, config_entry, async_add_entities):
@callback
def update_controller():
"""Update the values of the controller."""
update_items(controller, async_add_entities, tracked)
add_entities(controller, async_add_entities, tracked)
controller.listeners.append(
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
def update_items(controller, async_add_entities, tracked):
"""Update tracked device state from the controller."""
def add_entities(controller, async_add_entities, tracked):
"""Add new tracker entities from the controller."""
new_tracked = []
for items, tracker_class in (
@ -109,8 +109,6 @@ def update_items(controller, async_add_entities, tracked):
for item_id in items:
if item_id in tracked:
if tracked[item_id].enabled:
tracked[item_id].async_schedule_update_ha_state()
continue
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)
class UniFiClientTracker(ScannerEntity):
class UniFiClientTracker(UniFiClient, ScannerEntity):
"""Representation of a network client."""
def __init__(self, client, controller):
"""Set up tracked client."""
self.client = client
self.controller = controller
self.is_wired = self.client.mac not in controller.wireless_clients
self.wired_bug = None
super().__init__(client, controller)
self.wired_bug = None
if self.is_wired != self.client.is_wired:
self.wired_bug = dt_util.utcnow() - self.controller.option_detection_time
@ -151,26 +147,6 @@ class UniFiClientTracker(ScannerEntity):
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
def is_connected(self):
"""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 SOURCE_TYPE_ROUTER
@property
def name(self) -> str:
"""Return the name of the client."""
return self.client.name or self.client.hostname
@property
def unique_id(self) -> str:
"""Return a unique identifier for this client."""
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
def device_state_attributes(self):
"""Return the client state attributes."""
@ -239,6 +200,7 @@ class UniFiDeviceTracker(ScannerEntity):
"""Set up tracked device."""
self.device = device
self.controller = controller
self.listeners = []
@property
def entity_registry_enabled_default(self):
@ -250,17 +212,26 @@ class UniFiDeviceTracker(ScannerEntity):
async def async_added_to_hass(self):
"""Subscribe to device events."""
LOGGER.debug("New UniFi device tracker %s (%s)", self.name, self.device.mac)
async def async_update(self):
"""Synchronize state with controller."""
await self.controller.request_update()
LOGGER.debug(
"Updating UniFi tracked device %s\n%s",
self.entity_id,
pformat(self.device.raw),
self.device.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 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
def is_connected(self):
"""Return true if the device is connected to the network."""
@ -325,3 +296,8 @@ class UniFiDeviceTracker(ScannerEntity):
attributes["upgradable"] = self.device.upgradable
return attributes
@property
def should_poll(self):
"""No polling needed."""
return False

View File

@ -3,8 +3,12 @@
"name": "Ubiquiti UniFi",
"config_flow": true,
"documentation": "https://www.home-assistant.io/integrations/unifi",
"requirements": ["aiounifi==11"],
"requirements": [
"aiounifi==12"
],
"dependencies": [],
"codeowners": ["@kane610"],
"codeowners": [
"@kane610"
],
"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.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.entity import Entity
from homeassistant.helpers.entity_registry import DISABLED_CONFIG_ENTRY
from .unifi_client import UniFiClient
LOGGER = logging.getLogger(__name__)
ATTR_RECEIVING = "receiving"
@ -29,7 +29,7 @@ async def async_setup_entry(hass, config_entry, async_add_entities):
@callback
def update_controller():
"""Update the values of the controller."""
update_items(controller, async_add_entities, sensors)
add_entities(controller, async_add_entities, sensors)
controller.listeners.append(
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
def update_items(controller, async_add_entities, sensors):
"""Update sensors from the controller."""
def add_entities(controller, async_add_entities, sensors):
"""Add new sensor entities from the controller."""
new_sensors = []
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}"
if item_id in sensors:
sensor = sensors[item_id]
if sensor.enabled:
sensor.async_schedule_update_ha_state()
continue
sensors[item_id] = sensor_class(
@ -87,14 +84,8 @@ def update_items(controller, async_add_entities, sensors):
async_add_entities(new_sensors)
class UniFiBandwidthSensor(Entity):
"""UniFi Bandwidth sensor base class."""
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
class UniFiRxBandwidthSensor(UniFiClient):
"""Receiving bandwidth sensor."""
@property
def entity_registry_enabled_default(self):
@ -103,37 +94,6 @@ class UniFiBandwidthSensor(Entity):
return True
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
def state(self):
"""Return the state of the sensor."""
@ -153,7 +113,7 @@ class UniFiRxBandwidthSensor(UniFiBandwidthSensor):
return f"rx-{self.client.mac}"
class UniFiTxBandwidthSensor(UniFiBandwidthSensor):
class UniFiTxBandwidthSensor(UniFiRxBandwidthSensor):
"""Transmitting bandwidth sensor."""
@property

View File

@ -1,15 +1,15 @@
"""Support for devices connected to UniFi POE."""
import logging
from pprint import pformat
from homeassistant.components.switch import SwitchDevice
from homeassistant.components.unifi.config_flow import get_controller_from_config_entry
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 .unifi_client import UniFiClient
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):
"""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)
@ -55,7 +55,7 @@ async def async_setup_entry(hass, config_entry, async_add_entities):
@callback
def update_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(
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
def update_items(controller, async_add_entities, switches, switches_off):
"""Update POE port state from the controller."""
def add_entities(controller, async_add_entities, switches, switches_off):
"""Add new switch entities from the controller."""
new_switches = []
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}"
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
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}"
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
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)
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):
"""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.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."""
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()
if state is None:
@ -198,16 +163,6 @@ class UniFiPOEClientSwitch(UniFiClient, SwitchDevice, RestoreEntity):
if not self.client.sw_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
def unique_id(self):
"""Return a unique identifier for this switch."""
@ -267,10 +222,6 @@ class UniFiPOEClientSwitch(UniFiClient, SwitchDevice, RestoreEntity):
class UniFiBlockClientSwitch(UniFiClient, SwitchDevice):
"""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
def unique_id(self):
"""Return a unique identifier for this switch."""
@ -281,11 +232,6 @@ class UniFiBlockClientSwitch(UniFiClient, SwitchDevice):
"""Return true if client is allowed to connect."""
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):
"""Turn on connectivity for client."""
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
# homeassistant.components.unifi
aiounifi==11
aiounifi==12
# homeassistant.components.wwlln
aiowwlln==2.0.2

View File

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

View File

@ -111,9 +111,12 @@ async def setup_unifi_integration(
return mock_client_all_responses.popleft()
return {}
# "aiounifi.Controller.start_websocket", return_value=True
with patch("aiounifi.Controller.login", return_value=True), patch(
"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.async_block_till_done()
@ -233,47 +236,28 @@ async def test_reset_after_successful_setup(hass):
assert len(controller.listeners) == 0
async def test_failed_update_failed_login(hass):
"""Running update can handle a failed login."""
async def test_wireless_client_event_calls_update_wireless_devices(hass):
"""Call update_wireless_devices method when receiving wireless client event."""
controller = await setup_unifi_integration(hass)
with patch.object(
controller.api.clients, "update", side_effect=aiounifi.LoginRequired
), patch.object(controller.api, "login", side_effect=aiounifi.AiounifiException):
await controller.async_update()
await hass.async_block_till_done()
with patch(
"homeassistant.components.unifi.controller.UniFiController.update_wireless_clients",
return_value=None,
) as wireless_clients_mock:
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
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
assert wireless_clients_mock.assert_called_once
async def test_get_controller(hass):

View File

@ -2,6 +2,8 @@
from copy import copy
from datetime import timedelta
from aiounifi.controller import SIGNAL_CONNECTION_STATE
from aiounifi.websocket import STATE_DISCONNECTED, STATE_RUNNING
from asynctest import patch
from homeassistant import config_entries
@ -136,11 +138,12 @@ async def test_tracked_devices(hass):
client_1_copy = copy(CLIENT_1)
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["last_seen"] = dt_util.as_timestamp(dt_util.utcnow())
controller.mock_client_responses.append([client_1_copy])
controller.mock_device_responses.append([device_1_copy])
await controller.async_update()
event = {"meta": {"message": "device:sync"}, "data": [device_1_copy]}
controller.api.message_handler(event)
await hass.async_block_till_done()
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")
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["disabled"] = True
controller.mock_client_responses.append({})
controller.mock_device_responses.append([device_1_copy])
await controller.async_update()
event = {"meta": {"message": "device:sync"}, "data": [device_1_copy]}
controller.api.message_handler(event)
await hass.async_block_till_done()
device_1 = hass.states.get("device_tracker.device_1")
assert device_1.state == STATE_UNAVAILABLE
# Don't track wired clients nor devices
controller.config_entry.add_update_listener(controller.async_options_updated)
hass.config_entries.async_update_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["last_seen"] = dt_util.as_timestamp(dt_util.utcnow())
controller.mock_client_responses.append([client_1_client])
controller.mock_device_responses.append({})
await controller.async_update()
event = {"meta": {"message": "sta:sync"}, "data": [client_1_client]}
controller.api.message_handler(event)
await hass.async_block_till_done()
client_1 = hass.states.get("device_tracker.client_1")
@ -207,9 +232,8 @@ async def test_wireless_client_go_wired_issue(hass):
"utcnow",
return_value=(dt_util.utcnow() + timedelta(minutes=5)),
):
controller.mock_client_responses.append([client_1_client])
controller.mock_device_responses.append({})
await controller.async_update()
event = {"meta": {"message": "sta:sync"}, "data": [client_1_client]}
controller.api.message_handler(event)
await hass.async_block_till_done()
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["last_seen"] = dt_util.as_timestamp(dt_util.utcnow())
controller.mock_client_responses.append([client_1_client])
controller.mock_device_responses.append({})
await controller.async_update()
event = {"meta": {"message": "sta:sync"}, "data": [client_1_client]}
controller.api.message_handler(event)
await hass.async_block_till_done()
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.setup import async_setup_component
from .test_controller import setup_unifi_integration
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):
"""Test that configured options for a host are loaded via config entry."""
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)
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",
}
await setup_unifi_integration(hass)
assert hass.data[unifi.DOMAIN]
async def test_controller_fail_setup(hass):
"""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:
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] == {}
@ -140,33 +90,8 @@ async def test_controller_no_mac(hass):
async def test_unload_entry(hass):
"""Test being able to unload an entry."""
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)
controller = await setup_unifi_integration(hass)
assert hass.data[unifi.DOMAIN]
with patch.object(unifi, "UniFiController") as mock_controller, patch(
"homeassistant.helpers.device_registry.async_get_registry",
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] == {}
assert await unifi.async_unload_entry(hass, controller.config_entry)
assert not hass.data[unifi.DOMAIN]

View File

@ -90,8 +90,8 @@ async def test_sensors(hass):
clients[1]["rx_bytes"] = 2345000000
clients[1]["tx_bytes"] = 6789000000
controller.mock_client_responses.append(clients)
await controller.async_update()
event = {"meta": {"message": "sta:sync"}, "data": clients}
controller.api.message_handler(event)
await hass.async_block_till_done()
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(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
await hass.services.async_call(
"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 controller.mock_requests[3] == {
"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(
"switch", "turn_on", {"entity_id": "switch.block_client_1"}, blocking=True
)
assert len(controller.mock_requests) == 11
assert controller.mock_requests[7] == {
assert len(controller.mock_requests) == 5
assert controller.mock_requests[4] == {
"json": {"mac": "00:00:00:00:01:01", "cmd": "unblock-sta"},
"method": "post",
"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(hass.states.async_all()) == 2
controller.mock_client_responses.append([CLIENT_1, CLIENT_2])
controller.mock_device_responses.append([DEVICE_1])
controller.api.websocket._data = {
"meta": {"message": "sta:sync"},
"data": [CLIENT_2],
}
controller.api.session_handler("data")
# Calling a service will trigger the updates to run
await hass.services.async_call(
"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 controller.mock_requests[3] == {
"json": {
@ -360,7 +367,7 @@ async def test_new_client_discovered_on_poe_control(hass):
await hass.services.async_call(
"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] == {
"json": {
"port_overrides": [