diff --git a/homeassistant/components/unifi/__init__.py b/homeassistant/components/unifi/__init__.py index 65015b357a7..a21ae4ed508 100644 --- a/homeassistant/components/unifi/__init__.py +++ b/homeassistant/components/unifi/__init__.py @@ -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 diff --git a/homeassistant/components/unifi/controller.py b/homeassistant/components/unifi/controller.py index 826491f6ba6..27a0b6a668c 100644 --- a/homeassistant/components/unifi/controller.py +++ b/homeassistant/components/unifi/controller.py @@ -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: diff --git a/homeassistant/components/unifi/device_tracker.py b/homeassistant/components/unifi/device_tracker.py index 8b45a0f227b..859a37049b0 100644 --- a/homeassistant/components/unifi/device_tracker.py +++ b/homeassistant/components/unifi/device_tracker.py @@ -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 diff --git a/homeassistant/components/unifi/manifest.json b/homeassistant/components/unifi/manifest.json index e2bcd5b68a5..b4a4a5dab16 100644 --- a/homeassistant/components/unifi/manifest.json +++ b/homeassistant/components/unifi/manifest.json @@ -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" -} +} \ No newline at end of file diff --git a/homeassistant/components/unifi/sensor.py b/homeassistant/components/unifi/sensor.py index 9145fd8e00f..860ddf81d7d 100644 --- a/homeassistant/components/unifi/sensor.py +++ b/homeassistant/components/unifi/sensor.py @@ -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 diff --git a/homeassistant/components/unifi/switch.py b/homeassistant/components/unifi/switch.py index b1f62131eb4..be6002e886e 100644 --- a/homeassistant/components/unifi/switch.py +++ b/homeassistant/components/unifi/switch.py @@ -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) diff --git a/homeassistant/components/unifi/unifi_client.py b/homeassistant/components/unifi/unifi_client.py new file mode 100644 index 00000000000..2e18f55a57b --- /dev/null +++ b/homeassistant/components/unifi/unifi_client.py @@ -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 diff --git a/requirements_all.txt b/requirements_all.txt index 97532e06416..6997ea10869 100644 --- a/requirements_all.txt +++ b/requirements_all.txt @@ -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 diff --git a/requirements_test_all.txt b/requirements_test_all.txt index 99cc98edb91..62a755f5145 100644 --- a/requirements_test_all.txt +++ b/requirements_test_all.txt @@ -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 diff --git a/tests/components/unifi/test_controller.py b/tests/components/unifi/test_controller.py index a446bf914fb..e1b2b2355c4 100644 --- a/tests/components/unifi/test_controller.py +++ b/tests/components/unifi/test_controller.py @@ -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): diff --git a/tests/components/unifi/test_device_tracker.py b/tests/components/unifi/test_device_tracker.py index 429e685c574..86595fe43e2 100644 --- a/tests/components/unifi/test_device_tracker.py +++ b/tests/components/unifi/test_device_tracker.py @@ -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") diff --git a/tests/components/unifi/test_init.py b/tests/components/unifi/test_init.py index 1f5a3852e16..12f9c1bfd17 100644 --- a/tests/components/unifi/test_init.py +++ b/tests/components/unifi/test_init.py @@ -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] diff --git a/tests/components/unifi/test_sensor.py b/tests/components/unifi/test_sensor.py index 723d6871636..668b7a36ada 100644 --- a/tests/components/unifi/test_sensor.py +++ b/tests/components/unifi/test_sensor.py @@ -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") diff --git a/tests/components/unifi/test_switch.py b/tests/components/unifi/test_switch.py index cc4c41bcbfd..cd3d8785399 100644 --- a/tests/components/unifi/test_switch.py +++ b/tests/components/unifi/test_switch.py @@ -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": [