Rework UniFi wireless client "wired bug" logic (#89757)

This commit is contained in:
Robert Svensson 2023-03-24 10:06:09 +01:00 committed by GitHub
parent 7364e6ecb3
commit ee74e21541
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
6 changed files with 47 additions and 96 deletions

View File

@ -1,5 +1,7 @@
"""Integration to UniFi Network and its various features.""" """Integration to UniFi Network and its various features."""
from aiounifi.models.client import Client
from homeassistant.config_entries import ConfigEntry from homeassistant.config_entries import ConfigEntry
from homeassistant.const import EVENT_HOMEASSISTANT_STOP from homeassistant.const import EVENT_HOMEASSISTANT_STOP
from homeassistant.core import HomeAssistant, callback from homeassistant.core import HomeAssistant, callback
@ -91,33 +93,55 @@ def async_remove_poe_client_entities(
class UnifiWirelessClients: class UnifiWirelessClients:
"""Class to store clients known to be wireless. """Class to store clients known to be wireless.
This is needed since wireless devices going offline might get marked as wired by UniFi. This is needed since wireless devices going offline
might get marked as wired by UniFi.
""" """
def __init__(self, hass: HomeAssistant) -> None: def __init__(self, hass: HomeAssistant) -> None:
"""Set up client storage.""" """Set up client storage."""
self.hass = hass self.hass = hass
self.data: dict[str, dict[str, list[str]]] = {} self.data: dict[str, dict[str, list[str]] | list[str]] = {}
self.wireless_clients: set[str] = set()
self._store: Store = Store(hass, STORAGE_VERSION, STORAGE_KEY) self._store: Store = Store(hass, STORAGE_VERSION, STORAGE_KEY)
async def async_load(self) -> None: async def async_load(self) -> None:
"""Load data from file.""" """Load data from file."""
if (data := await self._store.async_load()) is not None: if (data := await self._store.async_load()) is not None:
self.data = data self.data = data
if "wireless_clients" not in data:
data["wireless_clients"] = [
obj_id
for config_entry in data
for obj_id in data[config_entry]["wireless_devices"]
]
self.wireless_clients.update(data["wireless_clients"])
@callback @callback
def get_data(self, config_entry: ConfigEntry) -> set[str]: def is_wireless(self, client: Client) -> bool:
"""Get data related to a specific controller.""" """Is client known to be wireless.
data = self.data.get(config_entry.entry_id, {"wireless_devices": []})
return set(data["wireless_devices"]) Store if client is wireless and not known.
"""
if not client.is_wired and client.mac not in self.wireless_clients:
self.wireless_clients.add(client.mac)
self._store.async_delay_save(self._data_to_save, SAVE_DELAY)
return client.mac in self.wireless_clients
@callback @callback
def update_data(self, data: set[str], config_entry: ConfigEntry) -> None: def update_clients(self, clients: set[Client]) -> None:
"""Update data and schedule to save to file.""" """Update data and schedule to save to file."""
self.data[config_entry.entry_id] = {"wireless_devices": list(data)} self.wireless_clients.update(
{client.mac for client in clients if not client.is_wired}
)
self._store.async_delay_save(self._data_to_save, SAVE_DELAY) self._store.async_delay_save(self._data_to_save, SAVE_DELAY)
@callback @callback
def _data_to_save(self) -> dict[str, dict[str, list[str]]]: def _data_to_save(self) -> dict[str, dict[str, list[str]] | list[str]]:
"""Return data of UniFi wireless clients to store in a file.""" """Return data of UniFi wireless clients to store in a file."""
self.data["wireless_clients"] = list(self.wireless_clients)
return self.data return self.data
def __contains__(self, obj_id: int | str) -> bool:
"""Validate membership of item ID."""
return obj_id in self.wireless_clients

View File

@ -10,8 +10,6 @@ from typing import Any
from aiohttp import CookieJar from aiohttp import CookieJar
import aiounifi import aiounifi
from aiounifi.interfaces.api_handlers import ItemEvent from aiounifi.interfaces.api_handlers import ItemEvent
from aiounifi.interfaces.messages import DATA_EVENT
from aiounifi.models.event import EventKey
from aiounifi.websocket import WebsocketSignal, WebsocketState from aiounifi.websocket import WebsocketSignal, WebsocketState
import async_timeout import async_timeout
@ -86,8 +84,7 @@ class UniFiController:
api.callback = self.async_unifi_signalling_callback api.callback = self.async_unifi_signalling_callback
self.available = True self.available = True
self.progress = None self.wireless_clients = hass.data[UNIFI_WIRELESS_CLIENTS]
self.wireless_clients = None
self.site_id: str = "" self.site_id: str = ""
self._site_name = None self._site_name = None
@ -247,15 +244,6 @@ class UniFiController:
else: else:
LOGGER.info("Connected to UniFi Network") LOGGER.info("Connected to UniFi Network")
elif signal == WebsocketSignal.DATA and DATA_EVENT in data:
for event in data[DATA_EVENT]:
if event.key in (
EventKey.WIRELESS_CLIENT_CONNECTED,
EventKey.WIRELESS_GUEST_CONNECTED,
):
self.update_wireless_clients()
break
@property @property
def signal_reachable(self) -> str: def signal_reachable(self) -> str:
"""Integration specific event to signal a change in connection status.""" """Integration specific event to signal a change in connection status."""
@ -271,22 +259,6 @@ class UniFiController:
"""Event specific per UniFi device tracker to signal new heartbeat missed.""" """Event specific per UniFi device tracker to signal new heartbeat missed."""
return "unifi-heartbeat-missed" return "unifi-heartbeat-missed"
def update_wireless_clients(self):
"""Update set of known to be wireless clients."""
new_wireless_clients = set()
for client_id in self.api.clients:
if (
client_id not in self.wireless_clients
and not self.api.clients[client_id].is_wired
):
new_wireless_clients.add(client_id)
if new_wireless_clients:
self.wireless_clients |= new_wireless_clients
unifi_wireless_clients = self.hass.data[UNIFI_WIRELESS_CLIENTS]
unifi_wireless_clients.update_data(self.wireless_clients, self.config_entry)
async def initialize(self): async def initialize(self):
"""Set up a UniFi Network instance.""" """Set up a UniFi Network instance."""
await self.api.initialize() await self.api.initialize()
@ -326,9 +298,7 @@ class UniFiController:
client.mac, client.mac,
) )
wireless_clients = self.hass.data[UNIFI_WIRELESS_CLIENTS] self.wireless_clients.update_clients(set(self.api.clients.values()))
self.wireless_clients = wireless_clients.get_data(self.config_entry)
self.update_wireless_clients()
self.config_entry.add_update_listener(self.async_config_entry_updated) self.config_entry.add_update_listener(self.async_config_entry_updated)

View File

@ -105,7 +105,7 @@ def async_client_is_connected_fn(controller: UniFiController, obj_id: str) -> bo
"""Check if device object is disabled.""" """Check if device object is disabled."""
client = controller.api.clients[obj_id] client = controller.api.clients[obj_id]
if client.is_wired != (obj_id not in controller.wireless_clients): if controller.wireless_clients.is_wireless(client) and client.is_wired:
if not controller.option_ignore_wired_bug: if not controller.option_ignore_wired_bug:
return False # Wired bug in action return False # Wired bug in action

View File

@ -45,17 +45,17 @@ from .entity import (
@callback @callback
def async_client_rx_value_fn(controller: UniFiController, client: Client) -> float: def async_client_rx_value_fn(controller: UniFiController, client: Client) -> float:
"""Calculate receiving data transfer value.""" """Calculate receiving data transfer value."""
if client.mac not in controller.wireless_clients: if controller.wireless_clients.is_wireless(client):
return client.wired_rx_bytes_r / 1000000 return client.rx_bytes_r / 1000000
return client.rx_bytes_r / 1000000 return client.wired_rx_bytes_r / 1000000
@callback @callback
def async_client_tx_value_fn(controller: UniFiController, client: Client) -> float: def async_client_tx_value_fn(controller: UniFiController, client: Client) -> float:
"""Calculate transmission data transfer value.""" """Calculate transmission data transfer value."""
if client.mac not in controller.wireless_clients: if controller.wireless_clients.is_wireless(client):
return client.wired_tx_bytes_r / 1000000 return client.tx_bytes_r / 1000000
return client.tx_bytes_r / 1000000 return client.wired_tx_bytes_r / 1000000
@callback @callback

View File

@ -6,8 +6,6 @@ from http import HTTPStatus
from unittest.mock import Mock, patch from unittest.mock import Mock, patch
import aiounifi import aiounifi
from aiounifi.models.event import EventKey
from aiounifi.models.message import MessageKey
from aiounifi.websocket import WebsocketState from aiounifi.websocket import WebsocketState
import pytest import pytest
@ -182,8 +180,8 @@ async def setup_unifi_integration(
config_entry.add_to_hass(hass) config_entry.add_to_hass(hass)
if known_wireless_clients: if known_wireless_clients:
hass.data[UNIFI_WIRELESS_CLIENTS].update_data( hass.data[UNIFI_WIRELESS_CLIENTS].wireless_clients.update(
known_wireless_clients, config_entry known_wireless_clients
) )
if aioclient_mock: if aioclient_mock:
@ -383,41 +381,6 @@ async def test_connection_state_signalling(
assert hass.states.get("device_tracker.client").state == "home" assert hass.states.get("device_tracker.client").state == "home"
async def test_wireless_client_event_calls_update_wireless_devices(
hass: HomeAssistant, aioclient_mock: AiohttpClientMocker, mock_unifi_websocket
) -> None:
"""Call update_wireless_devices method when receiving wireless client event."""
client_1_dict = {
"essid": "ssid",
"disabled": False,
"hostname": "client_1",
"ip": "10.0.0.4",
"is_wired": False,
"last_seen": dt_util.as_timestamp(dt_util.utcnow()),
"mac": "00:00:00:00:00:01",
}
await setup_unifi_integration(
hass,
aioclient_mock,
clients_response=[client_1_dict],
known_wireless_clients=(client_1_dict["mac"],),
)
with patch(
"homeassistant.components.unifi.controller.UniFiController.update_wireless_clients",
return_value=None,
) as wireless_clients_mock:
event = {
"datetime": "2020-01-20T19:37:04Z",
"user": "00:00:00:00:00:01",
"key": EventKey.WIRELESS_CLIENT_CONNECTED.value,
"msg": "User[11:22:33:44:55:66] has connected to WLAN",
"time": 1579549024893,
}
mock_unifi_websocket(message=MessageKey.EVENT, data=event)
assert wireless_clients_mock.assert_called_once
async def test_reconnect_mechanism( async def test_reconnect_mechanism(
hass: HomeAssistant, aioclient_mock: AiohttpClientMocker, mock_unifi_websocket hass: HomeAssistant, aioclient_mock: AiohttpClientMocker, mock_unifi_websocket
) -> None: ) -> None:

View File

@ -89,19 +89,13 @@ async def test_wireless_clients(
"is_wired": False, "is_wired": False,
"mac": "00:00:00:00:00:02", "mac": "00:00:00:00:00:02",
} }
config_entry = await setup_unifi_integration( await setup_unifi_integration(
hass, aioclient_mock, clients_response=[client_1, client_2] hass, aioclient_mock, clients_response=[client_1, client_2]
) )
await flush_store(hass.data[unifi.UNIFI_WIRELESS_CLIENTS]._store) await flush_store(hass.data[unifi.UNIFI_WIRELESS_CLIENTS]._store)
for mac in [ assert sorted(hass_storage[unifi.STORAGE_KEY]["data"]["wireless_clients"]) == [
"00:00:00:00:00:00", "00:00:00:00:00:00",
"00:00:00:00:00:01", "00:00:00:00:00:01",
"00:00:00:00:00:02", "00:00:00:00:00:02",
]: ]
assert (
mac
in hass_storage[unifi.STORAGE_KEY]["data"][config_entry.entry_id][
"wireless_devices"
]
)