Allow data sources to affect client tracker state after reconnecting to UniFi controller (#58269)

homeassistant/components/unifi/device_tracker.py

@@ -149,6 +149,7 @@ class UniFiClientTracker(UniFiClient, ScannerEntity):
         self.heartbeat_check = False
         self._is_connected = False
         self._controller_connection_state_changed = False
+        self._only_listen_to_event_source = False
 
         if client.last_seen:
             self._is_connected = (
@@ -191,11 +192,13 @@ class UniFiClientTracker(UniFiClient, ScannerEntity):
 
             if self.controller.available:
                 self.schedule_update = True
+                self._only_listen_to_event_source = False
 
             else:
                 self.controller.async_heartbeat(self.unique_id)
 
         elif self.client.last_updated == SOURCE_EVENT:
+            self._only_listen_to_event_source = True
             if (self.is_wired and self.client.event.event in WIRED_CONNECTION) or (
                 not self.is_wired and self.client.event.event in WIRELESS_CONNECTION
             ):
@@ -209,7 +212,7 @@ class UniFiClientTracker(UniFiClient, ScannerEntity):
                 self.schedule_update = True
 
         elif (
-            not self.client.event
+            not self._only_listen_to_event_source
             and self.client.last_updated == SOURCE_DATA
             and self.is_wired == self.client.is_wired
         ):

tests/components/unifi/test_device_tracker.py

@@ -68,7 +68,7 @@ async def test_tracked_wireless_clients(hass, aioclient_mock, mock_unifi_websock
     await hass.async_block_till_done()
 
     client_state = hass.states.get("device_tracker.client")
-    assert client_state.state == "home"
+    assert client_state.state == STATE_HOME
     assert client_state.attributes["ip"] == "10.0.0.1"
     assert client_state.attributes["mac"] == "00:00:00:00:00:01"
     assert client_state.attributes["hostname"] == "client"
@@ -112,6 +112,23 @@ async def test_tracked_wireless_clients(hass, aioclient_mock, mock_unifi_websock
 
     assert hass.states.get("device_tracker.client").state == STATE_NOT_HOME
 
+    # To limit false positives in client tracker
+    # data sources other than events are only used to update state
+    # until the first event has been received.
+    # This control will be reset if controller connection has been lost.
+
+    # New data doesn't change state
+
+    mock_unifi_websocket(
+        data={
+            "meta": {"message": MESSAGE_CLIENT},
+            "data": [client],
+        }
+    )
+    await hass.async_block_till_done()
+
+    assert hass.states.get("device_tracker.client").state == STATE_NOT_HOME
+
     # Connected event
 
     event = {
@@ -477,6 +494,91 @@ async def test_controller_state_change(hass, aioclient_mock, mock_unifi_websocke
     assert hass.states.get("device_tracker.device").state == STATE_HOME
 
 
+async def test_controller_state_change_client_to_listen_on_all_state_changes(
+    hass, aioclient_mock, mock_unifi_websocket
+):
+    """Verify entities state reflect on controller becoming unavailable."""
+    client = {
+        "ap_mac": "00:00:00:00:02:01",
+        "essid": "ssid",
+        "hostname": "client",
+        "ip": "10.0.0.1",
+        "is_wired": False,
+        "last_seen": dt_util.as_timestamp(dt_util.utcnow()),
+        "mac": "00:00:00:00:00:01",
+    }
+    config_entry = await setup_unifi_integration(
+        hass, aioclient_mock, clients_response=[client]
+    )
+    controller = hass.data[UNIFI_DOMAIN][config_entry.entry_id]
+
+    assert len(hass.states.async_entity_ids(TRACKER_DOMAIN)) == 1
+    assert hass.states.get("device_tracker.client").state == STATE_HOME
+
+    # Disconnected event
+
+    event = {
+        "user": client["mac"],
+        "ssid": client["essid"],
+        "hostname": client["hostname"],
+        "ap": client["ap_mac"],
+        "duration": 467,
+        "bytes": 459039,
+        "key": "EVT_WU_Disconnected",
+        "subsystem": "wlan",
+        "site_id": "name",
+        "time": 1587752927000,
+        "datetime": "2020-04-24T18:28:47Z",
+        "msg": f'User{[client["mac"]]} disconnected from "{client["essid"]}" (7m 47s connected, 448.28K bytes, last AP[{client["ap_mac"]}])',
+        "_id": "5ea32ff730c49e00f90dca1a",
+    }
+    mock_unifi_websocket(
+        data={
+            "meta": {"message": MESSAGE_EVENT},
+            "data": [event],
+        }
+    )
+    await hass.async_block_till_done()
+
+    assert hass.states.get("device_tracker.client").state == STATE_HOME
+
+    # Change time to mark client as away
+
+    new_time = dt_util.utcnow() + controller.option_detection_time
+    with patch("homeassistant.util.dt.utcnow", return_value=new_time):
+        async_fire_time_changed(hass, new_time)
+        await hass.async_block_till_done()
+
+    assert hass.states.get("device_tracker.client").state == STATE_NOT_HOME
+
+    # Controller unavailable
+    mock_unifi_websocket(state=STATE_DISCONNECTED)
+    await hass.async_block_till_done()
+
+    assert hass.states.get("device_tracker.client").state == STATE_UNAVAILABLE
+
+    # Controller available
+    mock_unifi_websocket(state=STATE_RUNNING)
+    await hass.async_block_till_done()
+
+    # To limit false positives in client tracker
+    # data sources other than events are only used to update state
+    # until the first event has been received.
+    # This control will be reset if controller connection has been lost.
+
+    # New data can change state
+
+    mock_unifi_websocket(
+        data={
+            "meta": {"message": MESSAGE_CLIENT},
+            "data": [client],
+        }
+    )
+    await hass.async_block_till_done()
+
+    assert hass.states.get("device_tracker.client").state == STATE_HOME
+
+
 async def test_option_track_clients(hass, aioclient_mock):
     """Test the tracking of clients can be turned off."""
     wireless_client = {