UniFi - Make devices proper push based (home-assistant#35152)

* Make devices proper push based

* Improve tests

* Bump dependency to v21
Update fix from home-assistant#35295 to use library

homeassistant/components/unifi/config_flow.py

@@ -239,7 +239,7 @@ async def async_step_device_tracker(self, user_input=None):
             | {
                 wlan["name"]
                 for ap in self.controller.api.devices.values()
-                for wlan in ap.raw.get("wlan_overrides", [])
+                for wlan in ap.wlan_overrides
             }
         )
         ssid_filter = {ssid: ssid for ssid in sorted(list(ssids))}

homeassistant/components/unifi/device_tracker.py

@@ -1,6 +1,9 @@
 """Track devices using UniFi controllers."""
+from datetime import timedelta
 import logging
 
+from aiounifi.api import SOURCE_DATA
+
 from homeassistant.components.device_tracker import DOMAIN
 from homeassistant.components.device_tracker.config_entry import ScannerEntity
 from homeassistant.components.device_tracker.const import SOURCE_TYPE_ROUTER
@@ -243,6 +246,9 @@ def __init__(self, device, controller):
         self.device = device
         super().__init__(controller)
 
+        self._is_connected = self.device.state == 1
+        self.cancel_scheduled_update = None
+
     @property
     def mac(self):
         """Return MAC of device."""
@@ -260,20 +266,34 @@ async def async_will_remove_from_hass(self) -> None:
 
     @callback
     def async_update_callback(self):
-        """Update the sensor's state."""
+        """Update the devices' state."""
+
+        @callback
+        def _no_heartbeat(now):
+            """No heart beat by device."""
+            self._is_connected = False
+            self.cancel_scheduled_update = None
+            self.async_write_ha_state()
+
+        if self.device.last_updated == SOURCE_DATA:
+            self._is_connected = True
+
+            if self.cancel_scheduled_update:
+                self.cancel_scheduled_update()
+
+            self.cancel_scheduled_update = async_track_point_in_utc_time(
+                self.hass,
+                _no_heartbeat,
+                dt_util.utcnow() + timedelta(seconds=self.device.next_interval + 10),
+            )
+
         LOGGER.debug("Updating device %s (%s)", self.entity_id, self.device.mac)
         self.async_write_ha_state()
 
     @property
     def is_connected(self):
         """Return true if the device is connected to the network."""
-        if self.device.state == 1 and (
-            dt_util.utcnow() - dt_util.utc_from_timestamp(float(self.device.last_seen))
-            < self.controller.option_detection_time
-        ):
-            return True
-
-        return False
+        return self._is_connected
 
     @property
     def source_type(self):
@@ -333,3 +353,8 @@ async def options_updated(self) -> None:
         """Config entry options are updated, remove entity if option is disabled."""
         if not self.controller.option_track_devices:
             await self.async_remove()
+
+    @property
+    def should_poll(self) -> bool:
+        """No polling needed."""
+        return False

homeassistant/components/unifi/manifest.json

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

requirements_all.txt

@@ -215,7 +215,7 @@ aiopylgtv==0.3.3
 aioswitcher==1.1.0
 
 # homeassistant.components.unifi
-aiounifi==20
+aiounifi==21
 
 # homeassistant.components.wwlln
 aiowwlln==2.0.2

requirements_test_all.txt

@@ -101,7 +101,7 @@ aiopylgtv==0.3.3
 aioswitcher==1.1.0
 
 # homeassistant.components.unifi
-aiounifi==20
+aiounifi==21
 
 # homeassistant.components.wwlln
 aiowwlln==2.0.2

tests/components/unifi/test_device_tracker.py

@@ -84,6 +84,7 @@
     "mac": "00:00:00:00:01:01",
     "model": "US16P150",
     "name": "device_1",
+    "next_interval": 20,
     "overheating": True,
     "state": 1,
     "type": "usw",
@@ -94,10 +95,11 @@
     "board_rev": 3,
     "device_id": "mock-id",
     "has_fan": True,
-    "ip": "10.0.1.1",
-    "mac": "00:00:00:00:01:01",
+    "ip": "10.0.1.2",
+    "mac": "00:00:00:00:01:02",
     "model": "US16P150",
-    "name": "device_1",
+    "name": "device_2",
+    "next_interval": 20,
     "state": 0,
     "type": "usw",
     "version": "4.0.42.10433",
@@ -206,7 +208,7 @@ async def test_tracked_wireless_clients(hass):
     # test wired bug
 
 
-async def test_tracked_devices(hass):
+async def test_tracked_clients(hass):
     """Test the update_items function with some clients."""
     client_4_copy = copy(CLIENT_4)
     client_4_copy["last_seen"] = dt_util.as_timestamp(dt_util.utcnow())
@@ -216,9 +218,9 @@ async def test_tracked_devices(hass):
         options={CONF_SSID_FILTER: ["ssid"]},
         clients_response=[CLIENT_1, CLIENT_2, CLIENT_3, CLIENT_5, client_4_copy],
         devices_response=[DEVICE_1, DEVICE_2],
-        known_wireless_clients=(CLIENT_4["mac"],),
+        known_wireless_clients=([CLIENT_4["mac"]]),
     )
-    assert len(hass.states.async_entity_ids(TRACKER_DOMAIN)) == 5
+    assert len(hass.states.async_entity_ids(TRACKER_DOMAIN)) == 6
 
     client_1 = hass.states.get("device_tracker.client_1")
     assert client_1 is not None
@@ -242,26 +244,54 @@ async def test_tracked_devices(hass):
     assert client_5 is not None
     assert client_5.state == "not_home"
 
-    device_1 = hass.states.get("device_tracker.device_1")
-    assert device_1 is not None
-    assert device_1.state == "not_home"
-
     # State change signalling works
     client_1_copy = copy(CLIENT_1)
     client_1_copy["last_seen"] = dt_util.as_timestamp(dt_util.utcnow())
     event = {"meta": {"message": MESSAGE_CLIENT}, "data": [client_1_copy]}
     controller.api.message_handler(event)
+
+    client_1 = hass.states.get("device_tracker.client_1")
+    assert client_1.state == "home"
+
+
+async def test_tracked_devices(hass):
+    """Test the update_items function with some devices."""
+    controller = await setup_unifi_integration(
+        hass, devices_response=[DEVICE_1, DEVICE_2],
+    )
+    assert len(hass.states.async_entity_ids(TRACKER_DOMAIN)) == 2
+
+    device_1 = hass.states.get("device_tracker.device_1")
+    assert device_1
+    assert device_1.state == "home"
+
+    device_2 = hass.states.get("device_tracker.device_2")
+    assert device_2
+    assert device_2.state == "not_home"
+
+    # State change signalling work
     device_1_copy = copy(DEVICE_1)
-    device_1_copy["last_seen"] = dt_util.as_timestamp(dt_util.utcnow())
+    device_1_copy["next_interval"] = 20
     event = {"meta": {"message": MESSAGE_DEVICE}, "data": [device_1_copy]}
     controller.api.message_handler(event)
+    device_2_copy = copy(DEVICE_2)
+    device_2_copy["next_interval"] = 50
+    event = {"meta": {"message": MESSAGE_DEVICE}, "data": [device_2_copy]}
+    controller.api.message_handler(event)
     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"
+    device_2 = hass.states.get("device_tracker.device_2")
+    assert device_2.state == "home"
+
+    async_fire_time_changed(hass, dt_util.utcnow() + timedelta(seconds=40))
+    await hass.async_block_till_done()
+
+    device_1 = hass.states.get("device_tracker.device_1")
+    assert device_1.state == "not_home"
+    device_2 = hass.states.get("device_tracker.device_2")
+    assert device_2.state == "home"
 
     # Disabled device is unavailable
     device_1_copy = copy(DEVICE_1)
@@ -330,7 +360,7 @@ async def test_controller_state_change(hass):
     assert client_1.state == "not_home"
 
     device_1 = hass.states.get("device_tracker.device_1")
-    assert device_1.state == "not_home"
+    assert device_1.state == "home"
 
 
 async def test_option_track_clients(hass):
@@ -648,7 +678,7 @@ async def test_dont_track_clients(hass):
 
     device_1 = hass.states.get("device_tracker.device_1")
     assert device_1 is not None
-    assert device_1.state == "not_home"
+    assert device_1.state == "home"
 
 
 async def test_dont_track_devices(hass):