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>

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
 
 

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(
+            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.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:

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):
+def add_entities(controller, async_add_entities, tracked):
-    """Update tracked device state from the controller."""
+    """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
+        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:
             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)
-
+        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
     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

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"
-}
+}

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):
+def add_entities(controller, async_add_entities, sensors):
-    """Update sensors from the controller."""
+    """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):
+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
     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

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):
+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 = []
     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)

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

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

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

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):
+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)
 
-    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):

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])
+    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()
 
     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({})
+    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()
 
     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])
+    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()
 
     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])
+        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()
 
         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])
+    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()
 
     client_1 = hass.states.get("device_tracker.client_1")

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(
+    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):
     """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(
+    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] == {}

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)
+    event = {"meta": {"message": "sta:sync"}, "data": clients}
-    await controller.async_update()
+    controller.api.message_handler(event)
     await hass.async_block_till_done()
 
     wireless_client_rx = hass.states.get("sensor.wireless_client_name_rx")

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 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"},
         "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.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
     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": [