From 26c5f892072fa4ce8d383960b6941e763b193709 Mon Sep 17 00:00:00 2001 From: Robert Svensson Date: Tue, 26 Oct 2021 22:04:16 +0200 Subject: [PATCH] Allow data sources to affect client tracker state after reconnecting to UniFi controller (#58269) --- .../components/unifi/device_tracker.py | 5 +- tests/components/unifi/test_device_tracker.py | 104 +++++++++++++++++- 2 files changed, 107 insertions(+), 2 deletions(-) diff --git a/homeassistant/components/unifi/device_tracker.py b/homeassistant/components/unifi/device_tracker.py index b6e14717484..ecb549e89ef 100644 --- a/homeassistant/components/unifi/device_tracker.py +++ b/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 ): diff --git a/tests/components/unifi/test_device_tracker.py b/tests/components/unifi/test_device_tracker.py index 4ba690dc444..5ab61ce21a7 100644 --- a/tests/components/unifi/test_device_tracker.py +++ b/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 = {