From 7c9891fff9f2ed00b730e9b561fb3a153dda96ba Mon Sep 17 00:00:00 2001 From: Robert Svensson Date: Tue, 5 Mar 2024 21:36:26 +0100 Subject: [PATCH] Break out UniFi config handling to own class (#111325) --- homeassistant/components/unifi/config_flow.py | 24 ++-- .../components/unifi/device_tracker.py | 22 ++-- homeassistant/components/unifi/hub/config.py | 123 ++++++++++++++++++ homeassistant/components/unifi/hub/hub.py | 114 +++------------- homeassistant/components/unifi/sensor.py | 20 +-- homeassistant/components/unifi/switch.py | 8 +- tests/components/unifi/conftest.py | 4 +- tests/components/unifi/test_hub.py | 20 +-- 8 files changed, 193 insertions(+), 142 deletions(-) create mode 100644 homeassistant/components/unifi/hub/config.py diff --git a/homeassistant/components/unifi/config_flow.py b/homeassistant/components/unifi/config_flow.py index 0342707c657..ae55d7ffcc9 100644 --- a/homeassistant/components/unifi/config_flow.py +++ b/homeassistant/components/unifi/config_flow.py @@ -260,7 +260,7 @@ class UnifiOptionsFlowHandler(OptionsFlow): if self.config_entry.entry_id not in self.hass.data[UNIFI_DOMAIN]: return self.async_abort(reason="integration_not_setup") self.hub = self.hass.data[UNIFI_DOMAIN][self.config_entry.entry_id] - self.options[CONF_BLOCK_CLIENT] = self.hub.option_block_clients + self.options[CONF_BLOCK_CLIENT] = self.hub.config.option_block_clients if self.show_advanced_options: return await self.async_step_configure_entity_sources() @@ -288,11 +288,11 @@ class UnifiOptionsFlowHandler(OptionsFlow): { vol.Optional( CONF_TRACK_CLIENTS, - default=self.hub.option_track_clients, + default=self.hub.config.option_track_clients, ): bool, vol.Optional( CONF_TRACK_DEVICES, - default=self.hub.option_track_devices, + default=self.hub.config.option_track_devices, ): bool, vol.Optional( CONF_BLOCK_CLIENT, default=self.options[CONF_BLOCK_CLIENT] @@ -361,7 +361,7 @@ class UnifiOptionsFlowHandler(OptionsFlow): ssid_filter = {ssid: ssid for ssid in sorted(ssids)} selected_ssids_to_filter = [ - ssid for ssid in self.hub.option_ssid_filter if ssid in ssid_filter + ssid for ssid in self.hub.config.option_ssid_filter if ssid in ssid_filter ] return self.async_show_form( @@ -370,26 +370,28 @@ class UnifiOptionsFlowHandler(OptionsFlow): { vol.Optional( CONF_TRACK_CLIENTS, - default=self.hub.option_track_clients, + default=self.hub.config.option_track_clients, ): bool, vol.Optional( CONF_TRACK_WIRED_CLIENTS, - default=self.hub.option_track_wired_clients, + default=self.hub.config.option_track_wired_clients, ): bool, vol.Optional( CONF_TRACK_DEVICES, - default=self.hub.option_track_devices, + default=self.hub.config.option_track_devices, ): bool, vol.Optional( CONF_SSID_FILTER, default=selected_ssids_to_filter ): cv.multi_select(ssid_filter), vol.Optional( CONF_DETECTION_TIME, - default=int(self.hub.option_detection_time.total_seconds()), + default=int( + self.hub.config.option_detection_time.total_seconds() + ), ): int, vol.Optional( CONF_IGNORE_WIRED_BUG, - default=self.hub.option_ignore_wired_bug, + default=self.hub.config.option_ignore_wired_bug, ): bool, } ), @@ -449,11 +451,11 @@ class UnifiOptionsFlowHandler(OptionsFlow): { vol.Optional( CONF_ALLOW_BANDWIDTH_SENSORS, - default=self.hub.option_allow_bandwidth_sensors, + default=self.hub.config.option_allow_bandwidth_sensors, ): bool, vol.Optional( CONF_ALLOW_UPTIME_SENSORS, - default=self.hub.option_allow_uptime_sensors, + default=self.hub.config.option_allow_uptime_sensors, ): bool, } ), diff --git a/homeassistant/components/unifi/device_tracker.py b/homeassistant/components/unifi/device_tracker.py index 81ba99c4ebe..cce41135e4c 100644 --- a/homeassistant/components/unifi/device_tracker.py +++ b/homeassistant/components/unifi/device_tracker.py @@ -82,21 +82,21 @@ WIRELESS_DISCONNECTION = ( @callback def async_client_allowed_fn(hub: UnifiHub, obj_id: str) -> bool: """Check if client is allowed.""" - if obj_id in hub.option_supported_clients: + if obj_id in hub.config.option_supported_clients: return True - if not hub.option_track_clients: + if not hub.config.option_track_clients: return False client = hub.api.clients[obj_id] if client.mac not in hub.wireless_clients: - if not hub.option_track_wired_clients: + if not hub.config.option_track_wired_clients: return False elif ( client.essid - and hub.option_ssid_filter - and client.essid not in hub.option_ssid_filter + and hub.config.option_ssid_filter + and client.essid not in hub.config.option_ssid_filter ): return False @@ -109,20 +109,20 @@ def async_client_is_connected_fn(hub: UnifiHub, obj_id: str) -> bool: client = hub.api.clients[obj_id] if hub.wireless_clients.is_wireless(client) and client.is_wired: - if not hub.option_ignore_wired_bug: + if not hub.config.option_ignore_wired_bug: return False # Wired bug in action if ( not client.is_wired and client.essid - and hub.option_ssid_filter - and client.essid not in hub.option_ssid_filter + and hub.config.option_ssid_filter + and client.essid not in hub.config.option_ssid_filter ): return False if ( dt_util.utcnow() - dt_util.utc_from_timestamp(client.last_seen or 0) - > hub.option_detection_time + > hub.config.option_detection_time ): return False @@ -161,7 +161,7 @@ ENTITY_DESCRIPTIONS: tuple[UnifiTrackerEntityDescription, ...] = ( + WIRELESS_CONNECTION + WIRELESS_DISCONNECTION ), - heartbeat_timedelta_fn=lambda hub, _: hub.option_detection_time, + heartbeat_timedelta_fn=lambda hub, _: hub.config.option_detection_time, is_connected_fn=async_client_is_connected_fn, name_fn=lambda client: client.name or client.hostname, object_fn=lambda api, obj_id: api.clients[obj_id], @@ -173,7 +173,7 @@ ENTITY_DESCRIPTIONS: tuple[UnifiTrackerEntityDescription, ...] = ( UnifiTrackerEntityDescription[Devices, Device]( key="Device scanner", has_entity_name=True, - allowed_fn=lambda hub, obj_id: hub.option_track_devices, + allowed_fn=lambda hub, obj_id: hub.config.option_track_devices, api_handler_fn=lambda api: api.devices, available_fn=async_device_available_fn, device_info_fn=lambda api, obj_id: None, diff --git a/homeassistant/components/unifi/hub/config.py b/homeassistant/components/unifi/hub/config.py new file mode 100644 index 00000000000..be2da7e47c9 --- /dev/null +++ b/homeassistant/components/unifi/hub/config.py @@ -0,0 +1,123 @@ +"""UniFi Network config entry abstraction.""" +from __future__ import annotations + +from dataclasses import dataclass +from datetime import timedelta +import ssl +from typing import Literal, Self + +from homeassistant.config_entries import ConfigEntry +from homeassistant.const import ( + CONF_HOST, + CONF_PASSWORD, + CONF_PORT, + CONF_USERNAME, + CONF_VERIFY_SSL, +) + +from ..const import ( + CONF_ALLOW_BANDWIDTH_SENSORS, + CONF_ALLOW_UPTIME_SENSORS, + CONF_BLOCK_CLIENT, + CONF_CLIENT_SOURCE, + CONF_DETECTION_TIME, + CONF_DPI_RESTRICTIONS, + CONF_IGNORE_WIRED_BUG, + CONF_SITE_ID, + CONF_SSID_FILTER, + CONF_TRACK_CLIENTS, + CONF_TRACK_DEVICES, + CONF_TRACK_WIRED_CLIENTS, + DEFAULT_ALLOW_BANDWIDTH_SENSORS, + DEFAULT_ALLOW_UPTIME_SENSORS, + DEFAULT_DETECTION_TIME, + DEFAULT_DPI_RESTRICTIONS, + DEFAULT_IGNORE_WIRED_BUG, + DEFAULT_TRACK_CLIENTS, + DEFAULT_TRACK_DEVICES, + DEFAULT_TRACK_WIRED_CLIENTS, +) + + +@dataclass +class UnifiConfig: + """Represent a UniFi config entry.""" + + entry: ConfigEntry + + host: str + port: int + username: str + password: str + site: str + ssl_context: ssl.SSLContext | Literal[False] + + option_supported_clients: list[str] + """Allow creating entities from clients.""" + + # Device tracker options + + option_track_clients: list[str] + """Config entry option to not track clients.""" + option_track_wired_clients: list[str] + """Config entry option to not track wired clients.""" + option_track_devices: bool + """Config entry option to not track devices.""" + option_ssid_filter: set[str] + """Config entry option listing what SSIDs are being used to track clients.""" + option_detection_time: timedelta + """Config entry option defining number of seconds from last seen to away""" + option_ignore_wired_bug: bool + """Config entry option to ignore wired bug.""" + + # Client control options + + option_block_clients: list[str] + """Config entry option with list of clients to control network access.""" + option_dpi_restrictions: bool + """Config entry option to control DPI restriction groups.""" + + # Statistics sensor options + + option_allow_bandwidth_sensors: bool + """Config entry option to allow bandwidth sensors.""" + option_allow_uptime_sensors: bool + """Config entry option to allow uptime sensors.""" + + @classmethod + def from_config_entry(cls, config_entry: ConfigEntry) -> Self: + """Create object from config entry.""" + config = config_entry.data + options = config_entry.options + return cls( + entry=config_entry, + host=config[CONF_HOST], + username=config[CONF_USERNAME], + password=config[CONF_PASSWORD], + port=config[CONF_PORT], + site=config[CONF_SITE_ID], + ssl_context=config.get(CONF_VERIFY_SSL, False), + option_supported_clients=options.get(CONF_CLIENT_SOURCE, []), + option_track_clients=options.get(CONF_TRACK_CLIENTS, DEFAULT_TRACK_CLIENTS), + option_track_wired_clients=options.get( + CONF_TRACK_WIRED_CLIENTS, DEFAULT_TRACK_WIRED_CLIENTS + ), + option_track_devices=options.get(CONF_TRACK_DEVICES, DEFAULT_TRACK_DEVICES), + option_ssid_filter=set(options.get(CONF_SSID_FILTER, [])), + option_detection_time=timedelta( + seconds=options.get(CONF_DETECTION_TIME, DEFAULT_DETECTION_TIME) + ), + option_ignore_wired_bug=options.get( + CONF_IGNORE_WIRED_BUG, DEFAULT_IGNORE_WIRED_BUG + ), + option_block_clients=options.get(CONF_BLOCK_CLIENT, []), + option_dpi_restrictions=options.get( + CONF_DPI_RESTRICTIONS, DEFAULT_DPI_RESTRICTIONS + ), + option_allow_bandwidth_sensors=options.get( + CONF_ALLOW_BANDWIDTH_SENSORS, DEFAULT_ALLOW_BANDWIDTH_SENSORS + ), + option_allow_uptime_sensors=options.get( + CONF_ALLOW_UPTIME_SENSORS, DEFAULT_ALLOW_UPTIME_SENSORS + ), + ) diff --git a/homeassistant/components/unifi/hub/hub.py b/homeassistant/components/unifi/hub/hub.py index 0188adf5c3f..6acfd1759f6 100644 --- a/homeassistant/components/unifi/hub/hub.py +++ b/homeassistant/components/unifi/hub/hub.py @@ -10,7 +10,7 @@ from aiounifi.interfaces.api_handlers import ItemEvent from aiounifi.models.device import DeviceSetPoePortModeRequest from homeassistant.config_entries import ConfigEntry -from homeassistant.const import CONF_HOST, Platform +from homeassistant.const import Platform from homeassistant.core import CALLBACK_TYPE, Event, HomeAssistant, callback from homeassistant.helpers import device_registry as dr, entity_registry as er from homeassistant.helpers.device_registry import ( @@ -29,31 +29,13 @@ import homeassistant.util.dt as dt_util from ..const import ( ATTR_MANUFACTURER, - CONF_ALLOW_BANDWIDTH_SENSORS, - CONF_ALLOW_UPTIME_SENSORS, - CONF_BLOCK_CLIENT, - CONF_CLIENT_SOURCE, - CONF_DETECTION_TIME, - CONF_DPI_RESTRICTIONS, - CONF_IGNORE_WIRED_BUG, CONF_SITE_ID, - CONF_SSID_FILTER, - CONF_TRACK_CLIENTS, - CONF_TRACK_DEVICES, - CONF_TRACK_WIRED_CLIENTS, - DEFAULT_ALLOW_BANDWIDTH_SENSORS, - DEFAULT_ALLOW_UPTIME_SENSORS, - DEFAULT_DETECTION_TIME, - DEFAULT_DPI_RESTRICTIONS, - DEFAULT_IGNORE_WIRED_BUG, - DEFAULT_TRACK_CLIENTS, - DEFAULT_TRACK_DEVICES, - DEFAULT_TRACK_WIRED_CLIENTS, DOMAIN as UNIFI_DOMAIN, PLATFORMS, UNIFI_WIRELESS_CLIENTS, ) from ..entity import UnifiEntity, UnifiEntityDescription +from .config import UnifiConfig from .websocket import UnifiWebsocket CHECK_HEARTBEAT_INTERVAL = timedelta(seconds=1) @@ -67,8 +49,8 @@ class UnifiHub: ) -> None: """Initialize the system.""" self.hass = hass - self.config_entry = config_entry self.api = api + self.config = UnifiConfig.from_config_entry(config_entry) self.websocket = UnifiWebsocket(hass, api, self.signal_reachable) self.wireless_clients = hass.data[UNIFI_WIRELESS_CLIENTS] @@ -79,72 +61,12 @@ class UnifiHub: self._cancel_heartbeat_check: CALLBACK_TYPE | None = None self._heartbeat_time: dict[str, datetime] = {} - self.load_config_entry_options() - self.entities: dict[str, str] = {} self.known_objects: set[tuple[str, str]] = set() self.poe_command_queue: dict[str, dict[int, str]] = {} self._cancel_poe_command: CALLBACK_TYPE | None = None - def load_config_entry_options(self) -> None: - """Store attributes to avoid property call overhead since they are called frequently.""" - options = self.config_entry.options - - # Allow creating entities from clients. - self.option_supported_clients: list[str] = options.get(CONF_CLIENT_SOURCE, []) - - # Device tracker options - - # Config entry option to not track clients. - self.option_track_clients = options.get( - CONF_TRACK_CLIENTS, DEFAULT_TRACK_CLIENTS - ) - # Config entry option to not track wired clients. - self.option_track_wired_clients = options.get( - CONF_TRACK_WIRED_CLIENTS, DEFAULT_TRACK_WIRED_CLIENTS - ) - # Config entry option to not track devices. - self.option_track_devices: bool = options.get( - CONF_TRACK_DEVICES, DEFAULT_TRACK_DEVICES - ) - # Config entry option listing what SSIDs are being used to track clients. - self.option_ssid_filter = set(options.get(CONF_SSID_FILTER, [])) - # Config entry option defining number of seconds from last seen to away - self.option_detection_time = timedelta( - seconds=options.get(CONF_DETECTION_TIME, DEFAULT_DETECTION_TIME) - ) - # Config entry option to ignore wired bug. - self.option_ignore_wired_bug = options.get( - CONF_IGNORE_WIRED_BUG, DEFAULT_IGNORE_WIRED_BUG - ) - - # Client control options - - # Config entry option with list of clients to control network access. - self.option_block_clients: list[str] = options.get(CONF_BLOCK_CLIENT, []) - # Config entry option to control DPI restriction groups. - self.option_dpi_restrictions: bool = options.get( - CONF_DPI_RESTRICTIONS, DEFAULT_DPI_RESTRICTIONS - ) - - # Statistics sensor options - - # Config entry option to allow bandwidth sensors. - self.option_allow_bandwidth_sensors: bool = options.get( - CONF_ALLOW_BANDWIDTH_SENSORS, DEFAULT_ALLOW_BANDWIDTH_SENSORS - ) - # Config entry option to allow uptime sensors. - self.option_allow_uptime_sensors: bool = options.get( - CONF_ALLOW_UPTIME_SENSORS, DEFAULT_ALLOW_UPTIME_SENSORS - ) - - @property - def host(self) -> str: - """Return the host of this hub.""" - host: str = self.config_entry.data[CONF_HOST] - return host - @property def available(self) -> bool: """Websocket connection state.""" @@ -219,7 +141,7 @@ class UnifiHub: partial(async_create_entity, description), ItemEvent.ADDED ) - self.config_entry.async_on_unload( + self.config.entry.async_on_unload( async_dispatcher_connect( self.hass, self.signal_options_update, @@ -232,12 +154,12 @@ class UnifiHub: @property def signal_reachable(self) -> str: """Integration specific event to signal a change in connection status.""" - return f"unifi-reachable-{self.config_entry.entry_id}" + return f"unifi-reachable-{self.config.entry.entry_id}" @property def signal_options_update(self) -> str: """Event specific per UniFi entry to signal new options.""" - return f"unifi-options-{self.config_entry.entry_id}" + return f"unifi-options-{self.config.entry.entry_id}" @property def signal_heartbeat_missed(self) -> str: @@ -248,25 +170,29 @@ class UnifiHub: """Set up a UniFi Network instance.""" await self.api.initialize() - assert self.config_entry.unique_id is not None - self.is_admin = self.api.sites[self.config_entry.unique_id].role == "admin" + assert self.config.entry.unique_id is not None + self.is_admin = self.api.sites[self.config.entry.unique_id].role == "admin" # Restore device tracker clients that are not a part of active clients list. macs: list[str] = [] entity_registry = er.async_get(self.hass) for entry in async_entries_for_config_entry( - entity_registry, self.config_entry.entry_id + entity_registry, self.config.entry.entry_id ): if entry.domain == Platform.DEVICE_TRACKER and "-" in entry.unique_id: macs.append(entry.unique_id.split("-", 1)[1]) - for mac in self.option_supported_clients + self.option_block_clients + macs: + for mac in ( + self.config.option_supported_clients + + self.config.option_block_clients + + macs + ): if mac not in self.api.clients and mac in self.api.clients_all: self.api.clients.process_raw([dict(self.api.clients_all[mac].raw)]) self.wireless_clients.update_clients(set(self.api.clients.values())) - self.config_entry.add_update_listener(self.async_config_entry_updated) + self.config.entry.add_update_listener(self.async_config_entry_updated) self._cancel_heartbeat_check = async_track_time_interval( self.hass, self._async_check_for_stale, CHECK_HEARTBEAT_INTERVAL @@ -328,7 +254,7 @@ class UnifiHub: @property def device_info(self) -> DeviceInfo: """UniFi Network device info.""" - assert self.config_entry.unique_id is not None + assert self.config.entry.unique_id is not None version: str | None = None if sysinfo := next(iter(self.api.system_information.values()), None): @@ -336,7 +262,7 @@ class UnifiHub: return DeviceInfo( entry_type=DeviceEntryType.SERVICE, - identifiers={(UNIFI_DOMAIN, self.config_entry.unique_id)}, + identifiers={(UNIFI_DOMAIN, self.config.entry.unique_id)}, manufacturer=ATTR_MANUFACTURER, model="UniFi Network Application", name="UniFi Network", @@ -348,7 +274,7 @@ class UnifiHub: """Update device registry.""" device_registry = dr.async_get(self.hass) return device_registry.async_get_or_create( - config_entry_id=self.config_entry.entry_id, **self.device_info + config_entry_id=self.config.entry.entry_id, **self.device_info ) @staticmethod @@ -362,7 +288,7 @@ class UnifiHub: """ if not (hub := hass.data[UNIFI_DOMAIN].get(config_entry.entry_id)): return - hub.load_config_entry_options() + hub.config = UnifiConfig.from_config_entry(config_entry) async_dispatcher_send(hass, hub.signal_options_update) @callback @@ -382,7 +308,7 @@ class UnifiHub: await self.websocket.stop_and_wait() unload_ok = await self.hass.config_entries.async_unload_platforms( - self.config_entry, PLATFORMS + self.config.entry, PLATFORMS ) if not unload_ok: diff --git a/homeassistant/components/unifi/sensor.py b/homeassistant/components/unifi/sensor.py index bf803b2c0ab..e83add10040 100644 --- a/homeassistant/components/unifi/sensor.py +++ b/homeassistant/components/unifi/sensor.py @@ -55,17 +55,17 @@ from .hub import UnifiHub @callback def async_bandwidth_sensor_allowed_fn(hub: UnifiHub, obj_id: str) -> bool: """Check if client is allowed.""" - if obj_id in hub.option_supported_clients: + if obj_id in hub.config.option_supported_clients: return True - return hub.option_allow_bandwidth_sensors + return hub.config.option_allow_bandwidth_sensors @callback def async_uptime_sensor_allowed_fn(hub: UnifiHub, obj_id: str) -> bool: """Check if client is allowed.""" - if obj_id in hub.option_supported_clients: + if obj_id in hub.config.option_supported_clients: return True - return hub.option_allow_uptime_sensors + return hub.config.option_allow_uptime_sensors @callback @@ -101,7 +101,7 @@ def async_wlan_client_value_fn(hub: UnifiHub, wlan: Wlan) -> int: for client in hub.api.clients.values() if client.essid == wlan.name and dt_util.utcnow() - dt_util.utc_from_timestamp(client.last_seen or 0) - < hub.option_detection_time + < hub.config.option_detection_time ] ) @@ -146,7 +146,7 @@ def async_client_is_connected_fn(hub: UnifiHub, obj_id: str) -> bool: if ( dt_util.utcnow() - dt_util.utc_from_timestamp(client.last_seen or 0) - > hub.option_detection_time + > hub.config.option_detection_time ): return False @@ -193,7 +193,7 @@ ENTITY_DESCRIPTIONS: tuple[UnifiSensorEntityDescription, ...] = ( is_connected_fn=async_client_is_connected_fn, name_fn=lambda _: "RX", object_fn=lambda api, obj_id: api.clients[obj_id], - supported_fn=lambda hub, _: hub.option_allow_bandwidth_sensors, + supported_fn=lambda hub, _: hub.config.option_allow_bandwidth_sensors, unique_id_fn=lambda hub, obj_id: f"rx-{obj_id}", value_fn=async_client_rx_value_fn, ), @@ -212,7 +212,7 @@ ENTITY_DESCRIPTIONS: tuple[UnifiSensorEntityDescription, ...] = ( is_connected_fn=async_client_is_connected_fn, name_fn=lambda _: "TX", object_fn=lambda api, obj_id: api.clients[obj_id], - supported_fn=lambda hub, _: hub.option_allow_bandwidth_sensors, + supported_fn=lambda hub, _: hub.config.option_allow_bandwidth_sensors, unique_id_fn=lambda hub, obj_id: f"tx-{obj_id}", value_fn=async_client_tx_value_fn, ), @@ -245,7 +245,7 @@ ENTITY_DESCRIPTIONS: tuple[UnifiSensorEntityDescription, ...] = ( device_info_fn=async_client_device_info_fn, name_fn=lambda client: "Uptime", object_fn=lambda api, obj_id: api.clients[obj_id], - supported_fn=lambda hub, _: hub.option_allow_uptime_sensors, + supported_fn=lambda hub, _: hub.config.option_allow_uptime_sensors, unique_id_fn=lambda hub, obj_id: f"uptime-{obj_id}", value_fn=async_client_uptime_value_fn, ), @@ -412,7 +412,7 @@ class UnifiSensorEntity(UnifiEntity[HandlerT, ApiItemT], SensorEntity): if description.is_connected_fn(self.hub, self._obj_id): self.hub.async_heartbeat( self._attr_unique_id, - dt_util.utcnow() + self.hub.option_detection_time, + dt_util.utcnow() + self.hub.config.option_detection_time, ) async def async_added_to_hass(self) -> None: diff --git a/homeassistant/components/unifi/switch.py b/homeassistant/components/unifi/switch.py index 6b2ce3c8b34..7e8a422ef23 100644 --- a/homeassistant/components/unifi/switch.py +++ b/homeassistant/components/unifi/switch.py @@ -64,9 +64,9 @@ CLIENT_UNBLOCKED = (EventKey.WIRED_CLIENT_UNBLOCKED, EventKey.WIRELESS_CLIENT_UN @callback def async_block_client_allowed_fn(hub: UnifiHub, obj_id: str) -> bool: """Check if client is allowed.""" - if obj_id in hub.option_supported_clients: + if obj_id in hub.config.option_supported_clients: return True - return obj_id in hub.option_block_clients + return obj_id in hub.config.option_block_clients @callback @@ -95,7 +95,7 @@ def async_dpi_group_device_info_fn(hub: UnifiHub, obj_id: str) -> DeviceInfo: @callback def async_port_forward_device_info_fn(hub: UnifiHub, obj_id: str) -> DeviceInfo: """Create device registry entry for port forward.""" - unique_id = hub.config_entry.unique_id + unique_id = hub.config.entry.unique_id assert unique_id is not None return DeviceInfo( entry_type=DeviceEntryType.SERVICE, @@ -203,7 +203,7 @@ ENTITY_DESCRIPTIONS: tuple[UnifiSwitchEntityDescription, ...] = ( key="DPI restriction", entity_category=EntityCategory.CONFIG, icon="mdi:network", - allowed_fn=lambda hub, obj_id: hub.option_dpi_restrictions, + allowed_fn=lambda hub, obj_id: hub.config.option_dpi_restrictions, api_handler_fn=lambda api: api.dpi_groups, available_fn=lambda hub, obj_id: hub.available, control_fn=async_dpi_group_control_fn, diff --git a/tests/components/unifi/conftest.py b/tests/components/unifi/conftest.py index 162f2b4d3aa..abc1592c4f5 100644 --- a/tests/components/unifi/conftest.py +++ b/tests/components/unifi/conftest.py @@ -45,10 +45,10 @@ class WebsocketStateManager(asyncio.Event): """ hub = self.hass.data[UNIFI_DOMAIN][DEFAULT_CONFIG_ENTRY_ID] self.aioclient_mock.get( - f"https://{hub.host}:1234", status=302 + f"https://{hub.config.host}:1234", status=302 ) # Check UniFi OS self.aioclient_mock.post( - f"https://{hub.host}:1234/api/login", + f"https://{hub.config.host}:1234/api/login", json={"data": "login successful", "meta": {"rc": "ok"}}, headers={"content-type": CONTENT_TYPE_JSON}, ) diff --git a/tests/components/unifi/test_hub.py b/tests/components/unifi/test_hub.py index 35b6e50cfd4..c1ea4f4ba0b 100644 --- a/tests/components/unifi/test_hub.py +++ b/tests/components/unifi/test_hub.py @@ -255,7 +255,7 @@ async def test_hub_setup( ) hub = hass.data[UNIFI_DOMAIN][config_entry.entry_id] - entry = hub.config_entry + entry = hub.config.entry assert len(forward_entry_setup.mock_calls) == len(PLATFORMS) assert forward_entry_setup.mock_calls[0][1] == (entry, BUTTON_DOMAIN) assert forward_entry_setup.mock_calls[1][1] == (entry, TRACKER_DOMAIN) @@ -263,17 +263,17 @@ async def test_hub_setup( assert forward_entry_setup.mock_calls[3][1] == (entry, SENSOR_DOMAIN) assert forward_entry_setup.mock_calls[4][1] == (entry, SWITCH_DOMAIN) - assert hub.host == ENTRY_CONFIG[CONF_HOST] + assert hub.config.host == ENTRY_CONFIG[CONF_HOST] assert hub.is_admin == (SITE[0]["role"] == "admin") - assert hub.option_allow_bandwidth_sensors == DEFAULT_ALLOW_BANDWIDTH_SENSORS - assert hub.option_allow_uptime_sensors == DEFAULT_ALLOW_UPTIME_SENSORS - assert isinstance(hub.option_block_clients, list) - assert hub.option_track_clients == DEFAULT_TRACK_CLIENTS - assert hub.option_track_devices == DEFAULT_TRACK_DEVICES - assert hub.option_track_wired_clients == DEFAULT_TRACK_WIRED_CLIENTS - assert hub.option_detection_time == timedelta(seconds=DEFAULT_DETECTION_TIME) - assert isinstance(hub.option_ssid_filter, set) + assert hub.config.option_allow_bandwidth_sensors == DEFAULT_ALLOW_BANDWIDTH_SENSORS + assert hub.config.option_allow_uptime_sensors == DEFAULT_ALLOW_UPTIME_SENSORS + assert isinstance(hub.config.option_block_clients, list) + assert hub.config.option_track_clients == DEFAULT_TRACK_CLIENTS + assert hub.config.option_track_devices == DEFAULT_TRACK_DEVICES + assert hub.config.option_track_wired_clients == DEFAULT_TRACK_WIRED_CLIENTS + assert hub.config.option_detection_time == timedelta(seconds=DEFAULT_DETECTION_TIME) + assert isinstance(hub.config.option_ssid_filter, set) assert hub.signal_reachable == "unifi-reachable-1" assert hub.signal_options_update == "unifi-options-1"