evlli/hass-core
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

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

Browse source
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
Show stats Download patch file Download diff file Expand all files Collapse all files

42
homeassistant/components/unifi/__init__.py
View file

@ -1,5 +1,7 @@
"""Integration to UniFi Network and its various features."""
from aiounifi.models.client import Client
from homeassistant.config_entries import ConfigEntry
from homeassistant.const import EVENT_HOMEASSISTANT_STOP
from homeassistant.core import HomeAssistant, callback
@ -91,33 +93,55 @@ def async_remove_poe_client_entities(
class UnifiWirelessClients:
"""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:
"""Set up client storage."""
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)
async def async_load(self) -> None:
"""Load data from file."""
if (data := await self._store.async_load()) is not None:
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
def get_data(self, config_entry: ConfigEntry) -> set[str]:
"""Get data related to a specific controller."""
data = self.data.get(config_entry.entry_id, {"wireless_devices": []})
return set(data["wireless_devices"])
def is_wireless(self, client: Client) -> bool:
"""Is client known to be wireless.
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
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."""
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)
@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."""
self.data["wireless_clients"] = list(self.wireless_clients)
return self.data
def __contains__(self, obj_id: int | str) -> bool:
"""Validate membership of item ID."""
return obj_id in self.wireless_clients

34
homeassistant/components/unifi/controller.py
View file

@ -10,8 +10,6 @@ from typing import Any
from aiohttp import CookieJar
import aiounifi
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
import async_timeout
@ -86,8 +84,7 @@ class UniFiController:
api.callback = self.async_unifi_signalling_callback
self.available = True
self.progress = None
self.wireless_clients = None
self.wireless_clients = hass.data[UNIFI_WIRELESS_CLIENTS]
self.site_id: str = ""
self._site_name = None
@ -247,15 +244,6 @@ class UniFiController:
else:
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
def signal_reachable(self) -> str:
"""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."""
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):
"""Set up a UniFi Network instance."""
await self.api.initialize()
@ -326,9 +298,7 @@ class UniFiController:
client.mac,
)
wireless_clients = self.hass.data[UNIFI_WIRELESS_CLIENTS]
self.wireless_clients = wireless_clients.get_data(self.config_entry)
self.update_wireless_clients()
self.wireless_clients.update_clients(set(self.api.clients.values()))
self.config_entry.add_update_listener(self.async_config_entry_updated)

2
homeassistant/components/unifi/device_tracker.py
View file

@ -105,7 +105,7 @@ def async_client_is_connected_fn(controller: UniFiController, obj_id: str) -> bo
"""Check if device object is disabled."""
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:
return False # Wired bug in action

8
homeassistant/components/unifi/sensor.py
View file

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

41
tests/components/unifi/test_controller.py
View file

@ -6,8 +6,6 @@ from http import HTTPStatus
from unittest.mock import Mock, patch
import aiounifi
from aiounifi.models.event import EventKey
from aiounifi.models.message import MessageKey
from aiounifi.websocket import WebsocketState
import pytest
@ -182,8 +180,8 @@ async def setup_unifi_integration(
config_entry.add_to_hass(hass)
if known_wireless_clients:
hass.data[UNIFI_WIRELESS_CLIENTS].update_data(
known_wireless_clients, config_entry
hass.data[UNIFI_WIRELESS_CLIENTS].wireless_clients.update(
known_wireless_clients
)
if aioclient_mock:
@ -383,41 +381,6 @@ async def test_connection_state_signalling(
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(
hass: HomeAssistant, aioclient_mock: AiohttpClientMocker, mock_unifi_websocket
) -> None:

10
tests/components/unifi/test_init.py
View file

@ -89,19 +89,13 @@ async def test_wireless_clients(
"is_wired": False,
"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]
)
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:01",
"00:00:00:00:00:02",
]:
assert (
mac
in hass_storage[unifi.STORAGE_KEY]["data"][config_entry.entry_id][
"wireless_devices"
]
)