Unifi rename controller to hub (#110976)

* Rename controller.py to hub.py

* Rename UniFiController to UnifiHub

* Rename controller instances into hub

* Rename controller to hub in tests

* Rename aiounifi Controller references to api

* Update strings

* Rename test_controller test_hub

* Narrow scope of test_remove_sensors
This commit is contained in:
Robert Svensson 2024-02-20 08:51:22 +01:00 committed by GitHub
parent 2b3f5319d6
commit 2f026ca963
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
25 changed files with 432 additions and 476 deletions

View file

@ -11,8 +11,8 @@ from homeassistant.helpers.storage import Store
from homeassistant.helpers.typing import ConfigType from homeassistant.helpers.typing import ConfigType
from .const import DOMAIN as UNIFI_DOMAIN, PLATFORMS, UNIFI_WIRELESS_CLIENTS from .const import DOMAIN as UNIFI_DOMAIN, PLATFORMS, UNIFI_WIRELESS_CLIENTS
from .controller import UniFiController, get_unifi_controller
from .errors import AuthenticationRequired, CannotConnect from .errors import AuthenticationRequired, CannotConnect
from .hub import UnifiHub, get_unifi_api
from .services import async_setup_services, async_unload_services from .services import async_setup_services, async_unload_services
SAVE_DELAY = 10 SAVE_DELAY = 10
@ -35,7 +35,7 @@ async def async_setup_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> b
hass.data.setdefault(UNIFI_DOMAIN, {}) hass.data.setdefault(UNIFI_DOMAIN, {})
try: try:
api = await get_unifi_controller(hass, config_entry.data) api = await get_unifi_api(hass, config_entry.data)
except CannotConnect as err: except CannotConnect as err:
raise ConfigEntryNotReady from err raise ConfigEntryNotReady from err
@ -43,20 +43,20 @@ async def async_setup_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> b
except AuthenticationRequired as err: except AuthenticationRequired as err:
raise ConfigEntryAuthFailed from err raise ConfigEntryAuthFailed from err
controller = UniFiController(hass, config_entry, api) hub = UnifiHub(hass, config_entry, api)
await controller.initialize() await hub.initialize()
hass.data[UNIFI_DOMAIN][config_entry.entry_id] = controller hass.data[UNIFI_DOMAIN][config_entry.entry_id] = hub
await hass.config_entries.async_forward_entry_setups(config_entry, PLATFORMS) await hass.config_entries.async_forward_entry_setups(config_entry, PLATFORMS)
controller.async_update_device_registry() hub.async_update_device_registry()
if len(hass.data[UNIFI_DOMAIN]) == 1: if len(hass.data[UNIFI_DOMAIN]) == 1:
async_setup_services(hass) async_setup_services(hass)
controller.start_websocket() hub.start_websocket()
config_entry.async_on_unload( config_entry.async_on_unload(
hass.bus.async_listen_once(EVENT_HOMEASSISTANT_STOP, controller.shutdown) hass.bus.async_listen_once(EVENT_HOMEASSISTANT_STOP, hub.shutdown)
) )
return True return True
@ -64,12 +64,12 @@ async def async_setup_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> b
async def async_unload_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> bool: async def async_unload_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> bool:
"""Unload a config entry.""" """Unload a config entry."""
controller: UniFiController = hass.data[UNIFI_DOMAIN].pop(config_entry.entry_id) hub: UnifiHub = hass.data[UNIFI_DOMAIN].pop(config_entry.entry_id)
if not hass.data[UNIFI_DOMAIN]: if not hass.data[UNIFI_DOMAIN]:
async_unload_services(hass) async_unload_services(hass)
return await controller.async_reset() return await hub.async_reset()
class UnifiWirelessClients: class UnifiWirelessClients:

View file

@ -30,7 +30,6 @@ from homeassistant.const import EntityCategory
from homeassistant.core import HomeAssistant, callback from homeassistant.core import HomeAssistant, callback
from homeassistant.helpers.entity_platform import AddEntitiesCallback from homeassistant.helpers.entity_platform import AddEntitiesCallback
from .controller import UniFiController
from .entity import ( from .entity import (
HandlerT, HandlerT,
UnifiEntity, UnifiEntity,
@ -38,6 +37,7 @@ from .entity import (
async_device_available_fn, async_device_available_fn,
async_device_device_info_fn, async_device_device_info_fn,
) )
from .hub import UnifiHub
@callback @callback
@ -79,7 +79,7 @@ ENTITY_DESCRIPTIONS: tuple[UnifiButtonEntityDescription, ...] = (
entity_category=EntityCategory.CONFIG, entity_category=EntityCategory.CONFIG,
has_entity_name=True, has_entity_name=True,
device_class=ButtonDeviceClass.RESTART, device_class=ButtonDeviceClass.RESTART,
allowed_fn=lambda controller, obj_id: True, allowed_fn=lambda hub, obj_id: True,
api_handler_fn=lambda api: api.devices, api_handler_fn=lambda api: api.devices,
available_fn=async_device_available_fn, available_fn=async_device_available_fn,
control_fn=async_restart_device_control_fn, control_fn=async_restart_device_control_fn,
@ -89,15 +89,15 @@ ENTITY_DESCRIPTIONS: tuple[UnifiButtonEntityDescription, ...] = (
name_fn=lambda _: "Restart", name_fn=lambda _: "Restart",
object_fn=lambda api, obj_id: api.devices[obj_id], object_fn=lambda api, obj_id: api.devices[obj_id],
should_poll=False, should_poll=False,
supported_fn=lambda controller, obj_id: True, supported_fn=lambda hub, obj_id: True,
unique_id_fn=lambda controller, obj_id: f"device_restart-{obj_id}", unique_id_fn=lambda hub, obj_id: f"device_restart-{obj_id}",
), ),
UnifiButtonEntityDescription[Ports, Port]( UnifiButtonEntityDescription[Ports, Port](
key="PoE power cycle", key="PoE power cycle",
entity_category=EntityCategory.CONFIG, entity_category=EntityCategory.CONFIG,
has_entity_name=True, has_entity_name=True,
device_class=ButtonDeviceClass.RESTART, device_class=ButtonDeviceClass.RESTART,
allowed_fn=lambda controller, obj_id: True, allowed_fn=lambda hub, obj_id: True,
api_handler_fn=lambda api: api.ports, api_handler_fn=lambda api: api.ports,
available_fn=async_device_available_fn, available_fn=async_device_available_fn,
control_fn=async_power_cycle_port_control_fn, control_fn=async_power_cycle_port_control_fn,
@ -107,8 +107,8 @@ ENTITY_DESCRIPTIONS: tuple[UnifiButtonEntityDescription, ...] = (
name_fn=lambda port: f"{port.name} Power Cycle", name_fn=lambda port: f"{port.name} Power Cycle",
object_fn=lambda api, obj_id: api.ports[obj_id], object_fn=lambda api, obj_id: api.ports[obj_id],
should_poll=False, should_poll=False,
supported_fn=lambda controller, obj_id: controller.api.ports[obj_id].port_poe, supported_fn=lambda hub, obj_id: hub.api.ports[obj_id].port_poe,
unique_id_fn=lambda controller, obj_id: f"power_cycle-{obj_id}", unique_id_fn=lambda hub, obj_id: f"power_cycle-{obj_id}",
), ),
) )
@ -119,7 +119,7 @@ async def async_setup_entry(
async_add_entities: AddEntitiesCallback, async_add_entities: AddEntitiesCallback,
) -> None: ) -> None:
"""Set up button platform for UniFi Network integration.""" """Set up button platform for UniFi Network integration."""
UniFiController.register_platform( UnifiHub.register_platform(
hass, hass,
config_entry, config_entry,
async_add_entities, async_add_entities,
@ -136,7 +136,7 @@ class UnifiButtonEntity(UnifiEntity[HandlerT, ApiItemT], ButtonEntity):
async def async_press(self) -> None: async def async_press(self) -> None:
"""Press the button.""" """Press the button."""
await self.entity_description.control_fn(self.controller.api, self._obj_id) await self.entity_description.control_fn(self.hub.api, self._obj_id)
@callback @callback
def async_update_state(self, event: ItemEvent, obj_id: str) -> None: def async_update_state(self, event: ItemEvent, obj_id: str) -> None:

View file

@ -47,8 +47,8 @@ from .const import (
DEFAULT_DPI_RESTRICTIONS, DEFAULT_DPI_RESTRICTIONS,
DOMAIN as UNIFI_DOMAIN, DOMAIN as UNIFI_DOMAIN,
) )
from .controller import UniFiController, get_unifi_controller
from .errors import AuthenticationRequired, CannotConnect from .errors import AuthenticationRequired, CannotConnect
from .hub import UnifiHub, get_unifi_api
DEFAULT_PORT = 443 DEFAULT_PORT = 443
DEFAULT_SITE_ID = "default" DEFAULT_SITE_ID = "default"
@ -99,11 +99,9 @@ class UnifiFlowHandler(config_entries.ConfigFlow, domain=UNIFI_DOMAIN):
} }
try: try:
controller = await get_unifi_controller( hub = await get_unifi_api(self.hass, MappingProxyType(self.config))
self.hass, MappingProxyType(self.config) await hub.sites.update()
) self.sites = hub.sites
await controller.sites.update()
self.sites = controller.sites
except AuthenticationRequired: except AuthenticationRequired:
errors["base"] = "faulty_credentials" errors["base"] = "faulty_credentials"
@ -160,11 +158,11 @@ class UnifiFlowHandler(config_entries.ConfigFlow, domain=UNIFI_DOMAIN):
abort_reason = "reauth_successful" abort_reason = "reauth_successful"
if config_entry: if config_entry:
controller: UniFiController | None = self.hass.data.get( hub: UnifiHub | None = self.hass.data.get(UNIFI_DOMAIN, {}).get(
UNIFI_DOMAIN, {} config_entry.entry_id
).get(config_entry.entry_id) )
if controller and controller.available: if hub and hub.available:
return self.async_abort(reason="already_configured") return self.async_abort(reason="already_configured")
return self.async_update_reload_and_abort( return self.async_update_reload_and_abort(
@ -240,7 +238,7 @@ class UnifiFlowHandler(config_entries.ConfigFlow, domain=UNIFI_DOMAIN):
class UnifiOptionsFlowHandler(config_entries.OptionsFlow): class UnifiOptionsFlowHandler(config_entries.OptionsFlow):
"""Handle Unifi Network options.""" """Handle Unifi Network options."""
controller: UniFiController hub: UnifiHub
def __init__(self, config_entry: config_entries.ConfigEntry) -> None: def __init__(self, config_entry: config_entries.ConfigEntry) -> None:
"""Initialize UniFi Network options flow.""" """Initialize UniFi Network options flow."""
@ -253,8 +251,8 @@ class UnifiOptionsFlowHandler(config_entries.OptionsFlow):
"""Manage the UniFi Network options.""" """Manage the UniFi Network options."""
if self.config_entry.entry_id not in self.hass.data[UNIFI_DOMAIN]: if self.config_entry.entry_id not in self.hass.data[UNIFI_DOMAIN]:
return self.async_abort(reason="integration_not_setup") return self.async_abort(reason="integration_not_setup")
self.controller = self.hass.data[UNIFI_DOMAIN][self.config_entry.entry_id] self.hub = self.hass.data[UNIFI_DOMAIN][self.config_entry.entry_id]
self.options[CONF_BLOCK_CLIENT] = self.controller.option_block_clients self.options[CONF_BLOCK_CLIENT] = self.hub.option_block_clients
if self.show_advanced_options: if self.show_advanced_options:
return await self.async_step_configure_entity_sources() return await self.async_step_configure_entity_sources()
@ -271,7 +269,7 @@ class UnifiOptionsFlowHandler(config_entries.OptionsFlow):
clients_to_block = {} clients_to_block = {}
for client in self.controller.api.clients.values(): for client in self.hub.api.clients.values():
clients_to_block[ clients_to_block[
client.mac client.mac
] = f"{client.name or client.hostname} ({client.mac})" ] = f"{client.name or client.hostname} ({client.mac})"
@ -282,11 +280,11 @@ class UnifiOptionsFlowHandler(config_entries.OptionsFlow):
{ {
vol.Optional( vol.Optional(
CONF_TRACK_CLIENTS, CONF_TRACK_CLIENTS,
default=self.controller.option_track_clients, default=self.hub.option_track_clients,
): bool, ): bool,
vol.Optional( vol.Optional(
CONF_TRACK_DEVICES, CONF_TRACK_DEVICES,
default=self.controller.option_track_devices, default=self.hub.option_track_devices,
): bool, ): bool,
vol.Optional( vol.Optional(
CONF_BLOCK_CLIENT, default=self.options[CONF_BLOCK_CLIENT] CONF_BLOCK_CLIENT, default=self.options[CONF_BLOCK_CLIENT]
@ -306,7 +304,7 @@ class UnifiOptionsFlowHandler(config_entries.OptionsFlow):
clients = { clients = {
client.mac: f"{client.name or client.hostname} ({client.mac})" client.mac: f"{client.name or client.hostname} ({client.mac})"
for client in self.controller.api.clients.values() for client in self.hub.api.clients.values()
} }
clients |= { clients |= {
mac: f"Unknown ({mac})" mac: f"Unknown ({mac})"
@ -338,16 +336,16 @@ class UnifiOptionsFlowHandler(config_entries.OptionsFlow):
return await self.async_step_client_control() return await self.async_step_client_control()
ssids = ( ssids = (
{wlan.name for wlan in self.controller.api.wlans.values()} {wlan.name for wlan in self.hub.api.wlans.values()}
| { | {
f"{wlan.name}{wlan.name_combine_suffix}" f"{wlan.name}{wlan.name_combine_suffix}"
for wlan in self.controller.api.wlans.values() for wlan in self.hub.api.wlans.values()
if not wlan.name_combine_enabled if not wlan.name_combine_enabled
and wlan.name_combine_suffix is not None and wlan.name_combine_suffix is not None
} }
| { | {
wlan["name"] wlan["name"]
for ap in self.controller.api.devices.values() for ap in self.hub.api.devices.values()
for wlan in ap.wlan_overrides for wlan in ap.wlan_overrides
if "name" in wlan if "name" in wlan
} }
@ -355,7 +353,7 @@ class UnifiOptionsFlowHandler(config_entries.OptionsFlow):
ssid_filter = {ssid: ssid for ssid in sorted(ssids)} ssid_filter = {ssid: ssid for ssid in sorted(ssids)}
selected_ssids_to_filter = [ selected_ssids_to_filter = [
ssid for ssid in self.controller.option_ssid_filter if ssid in ssid_filter ssid for ssid in self.hub.option_ssid_filter if ssid in ssid_filter
] ]
return self.async_show_form( return self.async_show_form(
@ -364,28 +362,26 @@ class UnifiOptionsFlowHandler(config_entries.OptionsFlow):
{ {
vol.Optional( vol.Optional(
CONF_TRACK_CLIENTS, CONF_TRACK_CLIENTS,
default=self.controller.option_track_clients, default=self.hub.option_track_clients,
): bool, ): bool,
vol.Optional( vol.Optional(
CONF_TRACK_WIRED_CLIENTS, CONF_TRACK_WIRED_CLIENTS,
default=self.controller.option_track_wired_clients, default=self.hub.option_track_wired_clients,
): bool, ): bool,
vol.Optional( vol.Optional(
CONF_TRACK_DEVICES, CONF_TRACK_DEVICES,
default=self.controller.option_track_devices, default=self.hub.option_track_devices,
): bool, ): bool,
vol.Optional( vol.Optional(
CONF_SSID_FILTER, default=selected_ssids_to_filter CONF_SSID_FILTER, default=selected_ssids_to_filter
): cv.multi_select(ssid_filter), ): cv.multi_select(ssid_filter),
vol.Optional( vol.Optional(
CONF_DETECTION_TIME, CONF_DETECTION_TIME,
default=int( default=int(self.hub.option_detection_time.total_seconds()),
self.controller.option_detection_time.total_seconds()
),
): int, ): int,
vol.Optional( vol.Optional(
CONF_IGNORE_WIRED_BUG, CONF_IGNORE_WIRED_BUG,
default=self.controller.option_ignore_wired_bug, default=self.hub.option_ignore_wired_bug,
): bool, ): bool,
} }
), ),
@ -402,7 +398,7 @@ class UnifiOptionsFlowHandler(config_entries.OptionsFlow):
clients_to_block = {} clients_to_block = {}
for client in self.controller.api.clients.values(): for client in self.hub.api.clients.values():
clients_to_block[ clients_to_block[
client.mac client.mac
] = f"{client.name or client.hostname} ({client.mac})" ] = f"{client.name or client.hostname} ({client.mac})"
@ -445,11 +441,11 @@ class UnifiOptionsFlowHandler(config_entries.OptionsFlow):
{ {
vol.Optional( vol.Optional(
CONF_ALLOW_BANDWIDTH_SENSORS, CONF_ALLOW_BANDWIDTH_SENSORS,
default=self.controller.option_allow_bandwidth_sensors, default=self.hub.option_allow_bandwidth_sensors,
): bool, ): bool,
vol.Optional( vol.Optional(
CONF_ALLOW_UPTIME_SENSORS, CONF_ALLOW_UPTIME_SENSORS,
default=self.controller.option_allow_uptime_sensors, default=self.hub.option_allow_uptime_sensors,
): bool, ): bool,
} }
), ),

View file

@ -25,13 +25,13 @@ from homeassistant.helpers.entity_platform import AddEntitiesCallback
import homeassistant.helpers.entity_registry as er import homeassistant.helpers.entity_registry as er
import homeassistant.util.dt as dt_util import homeassistant.util.dt as dt_util
from .controller import UNIFI_DOMAIN, UniFiController
from .entity import ( from .entity import (
HandlerT, HandlerT,
UnifiEntity, UnifiEntity,
UnifiEntityDescription, UnifiEntityDescription,
async_device_available_fn, async_device_available_fn,
) )
from .hub import UNIFI_DOMAIN, UnifiHub
LOGGER = logging.getLogger(__name__) LOGGER = logging.getLogger(__name__)
@ -79,23 +79,23 @@ WIRELESS_DISCONNECTION = (
@callback @callback
def async_client_allowed_fn(controller: UniFiController, obj_id: str) -> bool: def async_client_allowed_fn(hub: UnifiHub, obj_id: str) -> bool:
"""Check if client is allowed.""" """Check if client is allowed."""
if obj_id in controller.option_supported_clients: if obj_id in hub.option_supported_clients:
return True return True
if not controller.option_track_clients: if not hub.option_track_clients:
return False return False
client = controller.api.clients[obj_id] client = hub.api.clients[obj_id]
if client.mac not in controller.wireless_clients: if client.mac not in hub.wireless_clients:
if not controller.option_track_wired_clients: if not hub.option_track_wired_clients:
return False return False
elif ( elif (
client.essid client.essid
and controller.option_ssid_filter and hub.option_ssid_filter
and client.essid not in controller.option_ssid_filter and client.essid not in hub.option_ssid_filter
): ):
return False return False
@ -103,25 +103,25 @@ def async_client_allowed_fn(controller: UniFiController, obj_id: str) -> bool:
@callback @callback
def async_client_is_connected_fn(controller: UniFiController, obj_id: str) -> bool: def async_client_is_connected_fn(hub: UnifiHub, obj_id: str) -> bool:
"""Check if device object is disabled.""" """Check if device object is disabled."""
client = controller.api.clients[obj_id] client = hub.api.clients[obj_id]
if controller.wireless_clients.is_wireless(client) and client.is_wired: if hub.wireless_clients.is_wireless(client) and client.is_wired:
if not controller.option_ignore_wired_bug: if not hub.option_ignore_wired_bug:
return False # Wired bug in action return False # Wired bug in action
if ( if (
not client.is_wired not client.is_wired
and client.essid and client.essid
and controller.option_ssid_filter and hub.option_ssid_filter
and client.essid not in controller.option_ssid_filter and client.essid not in hub.option_ssid_filter
): ):
return False return False
if ( if (
dt_util.utcnow() - dt_util.utc_from_timestamp(client.last_seen or 0) dt_util.utcnow() - dt_util.utc_from_timestamp(client.last_seen or 0)
> controller.option_detection_time > hub.option_detection_time
): ):
return False return False
@ -129,11 +129,9 @@ def async_client_is_connected_fn(controller: UniFiController, obj_id: str) -> bo
@callback @callback
def async_device_heartbeat_timedelta_fn( def async_device_heartbeat_timedelta_fn(hub: UnifiHub, obj_id: str) -> timedelta:
controller: UniFiController, obj_id: str
) -> timedelta:
"""Check if device object is disabled.""" """Check if device object is disabled."""
device = controller.api.devices[obj_id] device = hub.api.devices[obj_id]
return timedelta(seconds=device.next_interval + 60) return timedelta(seconds=device.next_interval + 60)
@ -141,9 +139,9 @@ def async_device_heartbeat_timedelta_fn(
class UnifiEntityTrackerDescriptionMixin(Generic[HandlerT, ApiItemT]): class UnifiEntityTrackerDescriptionMixin(Generic[HandlerT, ApiItemT]):
"""Device tracker local functions.""" """Device tracker local functions."""
heartbeat_timedelta_fn: Callable[[UniFiController, str], timedelta] heartbeat_timedelta_fn: Callable[[UnifiHub, str], timedelta]
ip_address_fn: Callable[[aiounifi.Controller, str], str | None] ip_address_fn: Callable[[aiounifi.Controller, str], str | None]
is_connected_fn: Callable[[UniFiController, str], bool] is_connected_fn: Callable[[UnifiHub, str], bool]
hostname_fn: Callable[[aiounifi.Controller, str], str | None] hostname_fn: Callable[[aiounifi.Controller, str], str | None]
@ -161,7 +159,7 @@ ENTITY_DESCRIPTIONS: tuple[UnifiTrackerEntityDescription, ...] = (
has_entity_name=True, has_entity_name=True,
allowed_fn=async_client_allowed_fn, allowed_fn=async_client_allowed_fn,
api_handler_fn=lambda api: api.clients, api_handler_fn=lambda api: api.clients,
available_fn=lambda controller, obj_id: controller.available, available_fn=lambda hub, obj_id: hub.available,
device_info_fn=lambda api, obj_id: None, device_info_fn=lambda api, obj_id: None,
event_is_on=(WIRED_CONNECTION + WIRELESS_CONNECTION), event_is_on=(WIRED_CONNECTION + WIRELESS_CONNECTION),
event_to_subscribe=( event_to_subscribe=(
@ -170,20 +168,20 @@ ENTITY_DESCRIPTIONS: tuple[UnifiTrackerEntityDescription, ...] = (
+ WIRELESS_CONNECTION + WIRELESS_CONNECTION
+ WIRELESS_DISCONNECTION + WIRELESS_DISCONNECTION
), ),
heartbeat_timedelta_fn=lambda controller, _: controller.option_detection_time, heartbeat_timedelta_fn=lambda hub, _: hub.option_detection_time,
is_connected_fn=async_client_is_connected_fn, is_connected_fn=async_client_is_connected_fn,
name_fn=lambda client: client.name or client.hostname, name_fn=lambda client: client.name or client.hostname,
object_fn=lambda api, obj_id: api.clients[obj_id], object_fn=lambda api, obj_id: api.clients[obj_id],
should_poll=False, should_poll=False,
supported_fn=lambda controller, obj_id: True, supported_fn=lambda hub, obj_id: True,
unique_id_fn=lambda controller, obj_id: f"{controller.site}-{obj_id}", unique_id_fn=lambda hub, obj_id: f"{hub.site}-{obj_id}",
ip_address_fn=lambda api, obj_id: api.clients[obj_id].ip, ip_address_fn=lambda api, obj_id: api.clients[obj_id].ip,
hostname_fn=lambda api, obj_id: api.clients[obj_id].hostname, hostname_fn=lambda api, obj_id: api.clients[obj_id].hostname,
), ),
UnifiTrackerEntityDescription[Devices, Device]( UnifiTrackerEntityDescription[Devices, Device](
key="Device scanner", key="Device scanner",
has_entity_name=True, has_entity_name=True,
allowed_fn=lambda controller, obj_id: controller.option_track_devices, allowed_fn=lambda hub, obj_id: hub.option_track_devices,
api_handler_fn=lambda api: api.devices, api_handler_fn=lambda api: api.devices,
available_fn=async_device_available_fn, available_fn=async_device_available_fn,
device_info_fn=lambda api, obj_id: None, device_info_fn=lambda api, obj_id: None,
@ -194,8 +192,8 @@ ENTITY_DESCRIPTIONS: tuple[UnifiTrackerEntityDescription, ...] = (
name_fn=lambda device: device.name or device.model, name_fn=lambda device: device.name or device.model,
object_fn=lambda api, obj_id: api.devices[obj_id], object_fn=lambda api, obj_id: api.devices[obj_id],
should_poll=False, should_poll=False,
supported_fn=lambda controller, obj_id: True, supported_fn=lambda hub, obj_id: True,
unique_id_fn=lambda controller, obj_id: obj_id, unique_id_fn=lambda hub, obj_id: obj_id,
ip_address_fn=lambda api, obj_id: api.devices[obj_id].ip, ip_address_fn=lambda api, obj_id: api.devices[obj_id].ip,
hostname_fn=lambda api, obj_id: None, hostname_fn=lambda api, obj_id: None,
), ),
@ -208,21 +206,21 @@ def async_update_unique_id(hass: HomeAssistant, config_entry: ConfigEntry) -> No
Introduced with release 2023.12. Introduced with release 2023.12.
""" """
controller: UniFiController = hass.data[UNIFI_DOMAIN][config_entry.entry_id] hub: UnifiHub = hass.data[UNIFI_DOMAIN][config_entry.entry_id]
ent_reg = er.async_get(hass) ent_reg = er.async_get(hass)
@callback @callback
def update_unique_id(obj_id: str) -> None: def update_unique_id(obj_id: str) -> None:
"""Rework unique ID.""" """Rework unique ID."""
new_unique_id = f"{controller.site}-{obj_id}" new_unique_id = f"{hub.site}-{obj_id}"
if ent_reg.async_get_entity_id(DOMAIN, UNIFI_DOMAIN, new_unique_id): if ent_reg.async_get_entity_id(DOMAIN, UNIFI_DOMAIN, new_unique_id):
return return
unique_id = f"{obj_id}-{controller.site}" unique_id = f"{obj_id}-{hub.site}"
if entity_id := ent_reg.async_get_entity_id(DOMAIN, UNIFI_DOMAIN, unique_id): if entity_id := ent_reg.async_get_entity_id(DOMAIN, UNIFI_DOMAIN, unique_id):
ent_reg.async_update_entity(entity_id, new_unique_id=new_unique_id) ent_reg.async_update_entity(entity_id, new_unique_id=new_unique_id)
for obj_id in list(controller.api.clients) + list(controller.api.clients_all): for obj_id in list(hub.api.clients) + list(hub.api.clients_all):
update_unique_id(obj_id) update_unique_id(obj_id)
@ -233,7 +231,7 @@ async def async_setup_entry(
) -> None: ) -> None:
"""Set up device tracker for UniFi Network integration.""" """Set up device tracker for UniFi Network integration."""
async_update_unique_id(hass, config_entry) async_update_unique_id(hass, config_entry)
UniFiController.register_platform( UnifiHub.register_platform(
hass, config_entry, async_add_entities, UnifiScannerEntity, ENTITY_DESCRIPTIONS hass, config_entry, async_add_entities, UnifiScannerEntity, ENTITY_DESCRIPTIONS
) )
@ -256,12 +254,12 @@ class UnifiScannerEntity(UnifiEntity[HandlerT, ApiItemT], ScannerEntity):
description = self.entity_description description = self.entity_description
self._event_is_on = description.event_is_on or () self._event_is_on = description.event_is_on or ()
self._ignore_events = False self._ignore_events = False
self._is_connected = description.is_connected_fn(self.controller, self._obj_id) self._is_connected = description.is_connected_fn(self.hub, self._obj_id)
if self.is_connected: if self.is_connected:
self.controller.async_heartbeat( self.hub.async_heartbeat(
self.unique_id, self.unique_id,
dt_util.utcnow() dt_util.utcnow()
+ description.heartbeat_timedelta_fn(self.controller, self._obj_id), + description.heartbeat_timedelta_fn(self.hub, self._obj_id),
) )
@property @property
@ -272,12 +270,12 @@ class UnifiScannerEntity(UnifiEntity[HandlerT, ApiItemT], ScannerEntity):
@property @property
def hostname(self) -> str | None: def hostname(self) -> str | None:
"""Return hostname of the device.""" """Return hostname of the device."""
return self.entity_description.hostname_fn(self.controller.api, self._obj_id) return self.entity_description.hostname_fn(self.hub.api, self._obj_id)
@property @property
def ip_address(self) -> str | None: def ip_address(self) -> str | None:
"""Return the primary ip address of the device.""" """Return the primary ip address of the device."""
return self.entity_description.ip_address_fn(self.controller.api, self._obj_id) return self.entity_description.ip_address_fn(self.hub.api, self._obj_id)
@property @property
def mac_address(self) -> str: def mac_address(self) -> str:
@ -304,7 +302,7 @@ class UnifiScannerEntity(UnifiEntity[HandlerT, ApiItemT], ScannerEntity):
def async_update_state(self, event: ItemEvent, obj_id: str) -> None: def async_update_state(self, event: ItemEvent, obj_id: str) -> None:
"""Update entity state. """Update entity state.
Remove heartbeat check if controller state has changed Remove heartbeat check if hub connection state has changed
and entity is unavailable. and entity is unavailable.
Update is_connected. Update is_connected.
Schedule new heartbeat check if connected. Schedule new heartbeat check if connected.
@ -319,15 +317,15 @@ class UnifiScannerEntity(UnifiEntity[HandlerT, ApiItemT], ScannerEntity):
# From unifi.entity.async_signal_reachable_callback # From unifi.entity.async_signal_reachable_callback
# Controller connection state has changed and entity is unavailable # Controller connection state has changed and entity is unavailable
# Cancel heartbeat # Cancel heartbeat
self.controller.async_heartbeat(self.unique_id) self.hub.async_heartbeat(self.unique_id)
return return
if is_connected := description.is_connected_fn(self.controller, self._obj_id): if is_connected := description.is_connected_fn(self.hub, self._obj_id):
self._is_connected = is_connected self._is_connected = is_connected
self.controller.async_heartbeat( self.hub.async_heartbeat(
self.unique_id, self.unique_id,
dt_util.utcnow() dt_util.utcnow()
+ description.heartbeat_timedelta_fn(self.controller, self._obj_id), + description.heartbeat_timedelta_fn(self.hub, self._obj_id),
) )
@callback @callback
@ -337,17 +335,15 @@ class UnifiScannerEntity(UnifiEntity[HandlerT, ApiItemT], ScannerEntity):
return return
if event.key in self._event_is_on: if event.key in self._event_is_on:
self.controller.async_heartbeat(self.unique_id) self.hub.async_heartbeat(self.unique_id)
self._is_connected = True self._is_connected = True
self.async_write_ha_state() self.async_write_ha_state()
return return
self.controller.async_heartbeat( self.hub.async_heartbeat(
self.unique_id, self.unique_id,
dt_util.utcnow() dt_util.utcnow()
+ self.entity_description.heartbeat_timedelta_fn( + self.entity_description.heartbeat_timedelta_fn(self.hub, self._obj_id),
self.controller, self._obj_id
),
) )
async def async_added_to_hass(self) -> None: async def async_added_to_hass(self) -> None:
@ -356,7 +352,7 @@ class UnifiScannerEntity(UnifiEntity[HandlerT, ApiItemT], ScannerEntity):
self.async_on_remove( self.async_on_remove(
async_dispatcher_connect( async_dispatcher_connect(
self.hass, self.hass,
f"{self.controller.signal_heartbeat_missed}_{self.unique_id}", f"{self.hub.signal_heartbeat_missed}_{self.unique_id}",
self._make_disconnected, self._make_disconnected,
) )
) )
@ -364,7 +360,7 @@ class UnifiScannerEntity(UnifiEntity[HandlerT, ApiItemT], ScannerEntity):
async def async_will_remove_from_hass(self) -> None: async def async_will_remove_from_hass(self) -> None:
"""Disconnect object when removed.""" """Disconnect object when removed."""
await super().async_will_remove_from_hass() await super().async_will_remove_from_hass()
self.controller.async_heartbeat(self.unique_id) self.hub.async_heartbeat(self.unique_id)
@property @property
def extra_state_attributes(self) -> Mapping[str, Any] | None: def extra_state_attributes(self) -> Mapping[str, Any] | None:
@ -372,7 +368,7 @@ class UnifiScannerEntity(UnifiEntity[HandlerT, ApiItemT], ScannerEntity):
if self.entity_description.key != "Client device scanner": if self.entity_description.key != "Client device scanner":
return None return None
client = self.entity_description.object_fn(self.controller.api, self._obj_id) client = self.entity_description.object_fn(self.hub.api, self._obj_id)
raw = client.raw raw = client.raw
attributes_to_check = CLIENT_STATIC_ATTRIBUTES attributes_to_check = CLIENT_STATIC_ATTRIBUTES

View file

@ -12,7 +12,7 @@ from homeassistant.core import HomeAssistant, callback
from homeassistant.helpers.device_registry import format_mac from homeassistant.helpers.device_registry import format_mac
from .const import DOMAIN as UNIFI_DOMAIN from .const import DOMAIN as UNIFI_DOMAIN
from .controller import UniFiController from .hub import UnifiHub
TO_REDACT = {CONF_PASSWORD} TO_REDACT = {CONF_PASSWORD}
REDACT_CONFIG = {CONF_HOST, CONF_PASSWORD, CONF_USERNAME} REDACT_CONFIG = {CONF_HOST, CONF_PASSWORD, CONF_USERNAME}
@ -75,16 +75,16 @@ async def async_get_config_entry_diagnostics(
hass: HomeAssistant, config_entry: ConfigEntry hass: HomeAssistant, config_entry: ConfigEntry
) -> dict[str, Any]: ) -> dict[str, Any]:
"""Return diagnostics for a config entry.""" """Return diagnostics for a config entry."""
controller: UniFiController = hass.data[UNIFI_DOMAIN][config_entry.entry_id] hub: UnifiHub = hass.data[UNIFI_DOMAIN][config_entry.entry_id]
diag: dict[str, Any] = {} diag: dict[str, Any] = {}
macs_to_redact: dict[str, str] = {} macs_to_redact: dict[str, str] = {}
counter = 0 counter = 0
for mac in chain(controller.api.clients, controller.api.devices): for mac in chain(hub.api.clients, hub.api.devices):
macs_to_redact[mac] = format_mac(str(counter).zfill(12)) macs_to_redact[mac] = format_mac(str(counter).zfill(12))
counter += 1 counter += 1
for device in controller.api.devices.values(): for device in hub.api.devices.values():
for entry in device.raw.get("ethernet_table", []): for entry in device.raw.get("ethernet_table", []):
mac = entry.get("mac", "") mac = entry.get("mac", "")
if mac not in macs_to_redact: if mac not in macs_to_redact:
@ -94,26 +94,26 @@ async def async_get_config_entry_diagnostics(
diag["config"] = async_redact_data( diag["config"] = async_redact_data(
async_replace_dict_data(config_entry.as_dict(), macs_to_redact), REDACT_CONFIG async_replace_dict_data(config_entry.as_dict(), macs_to_redact), REDACT_CONFIG
) )
diag["role_is_admin"] = controller.is_admin diag["role_is_admin"] = hub.is_admin
diag["clients"] = { diag["clients"] = {
macs_to_redact[k]: async_redact_data( macs_to_redact[k]: async_redact_data(
async_replace_dict_data(v.raw, macs_to_redact), REDACT_CLIENTS async_replace_dict_data(v.raw, macs_to_redact), REDACT_CLIENTS
) )
for k, v in controller.api.clients.items() for k, v in hub.api.clients.items()
} }
diag["devices"] = { diag["devices"] = {
macs_to_redact[k]: async_redact_data( macs_to_redact[k]: async_redact_data(
async_replace_dict_data(v.raw, macs_to_redact), REDACT_DEVICES async_replace_dict_data(v.raw, macs_to_redact), REDACT_DEVICES
) )
for k, v in controller.api.devices.items() for k, v in hub.api.devices.items()
} }
diag["dpi_apps"] = {k: v.raw for k, v in controller.api.dpi_apps.items()} diag["dpi_apps"] = {k: v.raw for k, v in hub.api.dpi_apps.items()}
diag["dpi_groups"] = {k: v.raw for k, v in controller.api.dpi_groups.items()} diag["dpi_groups"] = {k: v.raw for k, v in hub.api.dpi_groups.items()}
diag["wlans"] = { diag["wlans"] = {
k: async_redact_data( k: async_redact_data(
async_replace_dict_data(v.raw, macs_to_redact), REDACT_WLANS async_replace_dict_data(v.raw, macs_to_redact), REDACT_WLANS
) )
for k, v in controller.api.wlans.items() for k, v in hub.api.wlans.items()
} }
return diag return diag

View file

@ -29,36 +29,36 @@ from homeassistant.helpers.entity import Entity, EntityDescription
from .const import ATTR_MANUFACTURER, DOMAIN from .const import ATTR_MANUFACTURER, DOMAIN
if TYPE_CHECKING: if TYPE_CHECKING:
from .controller import UniFiController from .hub import UnifiHub
HandlerT = TypeVar("HandlerT", bound=APIHandler) HandlerT = TypeVar("HandlerT", bound=APIHandler)
SubscriptionT = Callable[[CallbackType, ItemEvent], UnsubscribeType] SubscriptionT = Callable[[CallbackType, ItemEvent], UnsubscribeType]
@callback @callback
def async_device_available_fn(controller: UniFiController, obj_id: str) -> bool: def async_device_available_fn(hub: UnifiHub, obj_id: str) -> bool:
"""Check if device is available.""" """Check if device is available."""
if "_" in obj_id: # Sub device (outlet or port) if "_" in obj_id: # Sub device (outlet or port)
obj_id = obj_id.partition("_")[0] obj_id = obj_id.partition("_")[0]
device = controller.api.devices[obj_id] device = hub.api.devices[obj_id]
return controller.available and not device.disabled return hub.available and not device.disabled
@callback @callback
def async_wlan_available_fn(controller: UniFiController, obj_id: str) -> bool: def async_wlan_available_fn(hub: UnifiHub, obj_id: str) -> bool:
"""Check if WLAN is available.""" """Check if WLAN is available."""
wlan = controller.api.wlans[obj_id] wlan = hub.api.wlans[obj_id]
return controller.available and wlan.enabled return hub.available and wlan.enabled
@callback @callback
def async_device_device_info_fn(controller: UniFiController, obj_id: str) -> DeviceInfo: def async_device_device_info_fn(hub: UnifiHub, obj_id: str) -> DeviceInfo:
"""Create device registry entry for device.""" """Create device registry entry for device."""
if "_" in obj_id: # Sub device (outlet or port) if "_" in obj_id: # Sub device (outlet or port)
obj_id = obj_id.partition("_")[0] obj_id = obj_id.partition("_")[0]
device = controller.api.devices[obj_id] device = hub.api.devices[obj_id]
return DeviceInfo( return DeviceInfo(
connections={(CONNECTION_NETWORK_MAC, device.mac)}, connections={(CONNECTION_NETWORK_MAC, device.mac)},
manufacturer=ATTR_MANUFACTURER, manufacturer=ATTR_MANUFACTURER,
@ -70,9 +70,9 @@ def async_device_device_info_fn(controller: UniFiController, obj_id: str) -> Dev
@callback @callback
def async_wlan_device_info_fn(controller: UniFiController, obj_id: str) -> DeviceInfo: def async_wlan_device_info_fn(hub: UnifiHub, obj_id: str) -> DeviceInfo:
"""Create device registry entry for WLAN.""" """Create device registry entry for WLAN."""
wlan = controller.api.wlans[obj_id] wlan = hub.api.wlans[obj_id]
return DeviceInfo( return DeviceInfo(
entry_type=DeviceEntryType.SERVICE, entry_type=DeviceEntryType.SERVICE,
identifiers={(DOMAIN, wlan.id)}, identifiers={(DOMAIN, wlan.id)},
@ -83,9 +83,9 @@ def async_wlan_device_info_fn(controller: UniFiController, obj_id: str) -> Devic
@callback @callback
def async_client_device_info_fn(controller: UniFiController, obj_id: str) -> DeviceInfo: def async_client_device_info_fn(hub: UnifiHub, obj_id: str) -> DeviceInfo:
"""Create device registry entry for client.""" """Create device registry entry for client."""
client = controller.api.clients[obj_id] client = hub.api.clients[obj_id]
return DeviceInfo( return DeviceInfo(
connections={(CONNECTION_NETWORK_MAC, obj_id)}, connections={(CONNECTION_NETWORK_MAC, obj_id)},
default_manufacturer=client.oui, default_manufacturer=client.oui,
@ -97,17 +97,17 @@ def async_client_device_info_fn(controller: UniFiController, obj_id: str) -> Dev
class UnifiDescription(Generic[HandlerT, ApiItemT]): class UnifiDescription(Generic[HandlerT, ApiItemT]):
"""Validate and load entities from different UniFi handlers.""" """Validate and load entities from different UniFi handlers."""
allowed_fn: Callable[[UniFiController, str], bool] allowed_fn: Callable[[UnifiHub, str], bool]
api_handler_fn: Callable[[aiounifi.Controller], HandlerT] api_handler_fn: Callable[[aiounifi.Controller], HandlerT]
available_fn: Callable[[UniFiController, str], bool] available_fn: Callable[[UnifiHub, str], bool]
device_info_fn: Callable[[UniFiController, str], DeviceInfo | None] device_info_fn: Callable[[UnifiHub, str], DeviceInfo | None]
event_is_on: tuple[EventKey, ...] | None event_is_on: tuple[EventKey, ...] | None
event_to_subscribe: tuple[EventKey, ...] | None event_to_subscribe: tuple[EventKey, ...] | None
name_fn: Callable[[ApiItemT], str | None] name_fn: Callable[[ApiItemT], str | None]
object_fn: Callable[[aiounifi.Controller, str], ApiItemT] object_fn: Callable[[aiounifi.Controller, str], ApiItemT]
should_poll: bool should_poll: bool
supported_fn: Callable[[UniFiController, str], bool | None] supported_fn: Callable[[UnifiHub, str], bool | None]
unique_id_fn: Callable[[UniFiController, str], str] unique_id_fn: Callable[[UnifiHub, str], str]
@dataclass(frozen=True) @dataclass(frozen=True)
@ -124,36 +124,36 @@ class UnifiEntity(Entity, Generic[HandlerT, ApiItemT]):
def __init__( def __init__(
self, self,
obj_id: str, obj_id: str,
controller: UniFiController, hub: UnifiHub,
description: UnifiEntityDescription[HandlerT, ApiItemT], description: UnifiEntityDescription[HandlerT, ApiItemT],
) -> None: ) -> None:
"""Set up UniFi switch entity.""" """Set up UniFi switch entity."""
self._obj_id = obj_id self._obj_id = obj_id
self.controller = controller self.hub = hub
self.entity_description = description self.entity_description = description
controller.known_objects.add((description.key, obj_id)) hub.known_objects.add((description.key, obj_id))
self._removed = False self._removed = False
self._attr_available = description.available_fn(controller, obj_id) self._attr_available = description.available_fn(hub, obj_id)
self._attr_device_info = description.device_info_fn(controller, obj_id) self._attr_device_info = description.device_info_fn(hub, obj_id)
self._attr_should_poll = description.should_poll self._attr_should_poll = description.should_poll
self._attr_unique_id = description.unique_id_fn(controller, obj_id) self._attr_unique_id = description.unique_id_fn(hub, obj_id)
obj = description.object_fn(self.controller.api, obj_id) obj = description.object_fn(self.hub.api, obj_id)
self._attr_name = description.name_fn(obj) self._attr_name = description.name_fn(obj)
self.async_initiate_state() self.async_initiate_state()
async def async_added_to_hass(self) -> None: async def async_added_to_hass(self) -> None:
"""Register callbacks.""" """Register callbacks."""
description = self.entity_description description = self.entity_description
handler = description.api_handler_fn(self.controller.api) handler = description.api_handler_fn(self.hub.api)
@callback @callback
def unregister_object() -> None: def unregister_object() -> None:
"""Remove object ID from known_objects when unloaded.""" """Remove object ID from known_objects when unloaded."""
self.controller.known_objects.discard((description.key, self._obj_id)) self.hub.known_objects.discard((description.key, self._obj_id))
self.async_on_remove(unregister_object) self.async_on_remove(unregister_object)
@ -165,11 +165,11 @@ class UnifiEntity(Entity, Generic[HandlerT, ApiItemT]):
) )
) )
# State change from controller or websocket # State change from hub or websocket
self.async_on_remove( self.async_on_remove(
async_dispatcher_connect( async_dispatcher_connect(
self.hass, self.hass,
self.controller.signal_reachable, self.hub.signal_reachable,
self.async_signal_reachable_callback, self.async_signal_reachable_callback,
) )
) )
@ -178,7 +178,7 @@ class UnifiEntity(Entity, Generic[HandlerT, ApiItemT]):
self.async_on_remove( self.async_on_remove(
async_dispatcher_connect( async_dispatcher_connect(
self.hass, self.hass,
self.controller.signal_options_update, self.hub.signal_options_update,
self.async_signal_options_updated, self.async_signal_options_updated,
) )
) )
@ -186,7 +186,7 @@ class UnifiEntity(Entity, Generic[HandlerT, ApiItemT]):
# Subscribe to events if defined # Subscribe to events if defined
if description.event_to_subscribe is not None: if description.event_to_subscribe is not None:
self.async_on_remove( self.async_on_remove(
self.controller.api.events.subscribe( self.hub.api.events.subscribe(
self.async_event_callback, self.async_event_callback,
description.event_to_subscribe, description.event_to_subscribe,
) )
@ -200,22 +200,22 @@ class UnifiEntity(Entity, Generic[HandlerT, ApiItemT]):
return return
description = self.entity_description description = self.entity_description
if not description.supported_fn(self.controller, self._obj_id): if not description.supported_fn(self.hub, self._obj_id):
self.hass.async_create_task(self.remove_item({self._obj_id})) self.hass.async_create_task(self.remove_item({self._obj_id}))
return return
self._attr_available = description.available_fn(self.controller, self._obj_id) self._attr_available = description.available_fn(self.hub, self._obj_id)
self.async_update_state(event, obj_id) self.async_update_state(event, obj_id)
self.async_write_ha_state() self.async_write_ha_state()
@callback @callback
def async_signal_reachable_callback(self) -> None: def async_signal_reachable_callback(self) -> None:
"""Call when controller connection state change.""" """Call when hub connection state change."""
self.async_signalling_callback(ItemEvent.ADDED, self._obj_id) self.async_signalling_callback(ItemEvent.ADDED, self._obj_id)
async def async_signal_options_updated(self) -> None: async def async_signal_options_updated(self) -> None:
"""Config entry options are updated, remove entity if option is disabled.""" """Config entry options are updated, remove entity if option is disabled."""
if not self.entity_description.allowed_fn(self.controller, self._obj_id): if not self.entity_description.allowed_fn(self.hub, self._obj_id):
await self.remove_item({self._obj_id}) await self.remove_item({self._obj_id})
async def remove_item(self, keys: set) -> None: async def remove_item(self, keys: set) -> None:

View file

@ -15,7 +15,7 @@ class AuthenticationRequired(UnifiException):
class CannotConnect(UnifiException): class CannotConnect(UnifiException):
"""Unable to connect to the controller.""" """Unable to connect to UniFi Network."""
class LoginRequired(UnifiException): class LoginRequired(UnifiException):

View file

@ -80,7 +80,7 @@ CHECK_HEARTBEAT_INTERVAL = timedelta(seconds=1)
CHECK_WEBSOCKET_INTERVAL = timedelta(minutes=1) CHECK_WEBSOCKET_INTERVAL = timedelta(minutes=1)
class UniFiController: class UnifiHub:
"""Manages a single UniFi Network instance.""" """Manages a single UniFi Network instance."""
def __init__( def __init__(
@ -165,7 +165,7 @@ class UniFiController:
@property @property
def host(self) -> str: def host(self) -> str:
"""Return the host of this controller.""" """Return the host of this hub."""
host: str = self.config_entry.data[CONF_HOST] host: str = self.config_entry.data[CONF_HOST]
return host return host
@ -180,10 +180,10 @@ class UniFiController:
requires_admin: bool = False, requires_admin: bool = False,
) -> None: ) -> None:
"""Register platform for UniFi entity management.""" """Register platform for UniFi entity management."""
controller: UniFiController = hass.data[UNIFI_DOMAIN][config_entry.entry_id] hub: UnifiHub = hass.data[UNIFI_DOMAIN][config_entry.entry_id]
if requires_admin and not controller.is_admin: if requires_admin and not hub.is_admin:
return return
controller.register_platform_add_entities( hub.register_platform_add_entities(
entity_class, descriptions, async_add_entities entity_class, descriptions, async_add_entities
) )
@ -351,7 +351,7 @@ class UniFiController:
@property @property
def device_info(self) -> DeviceInfo: def device_info(self) -> DeviceInfo:
"""UniFi controller device info.""" """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 version: str | None = None
@ -384,10 +384,10 @@ class UniFiController:
If config entry is updated due to reauth flow If config entry is updated due to reauth flow
the entry might already have been reset and thus is not available. the entry might already have been reset and thus is not available.
""" """
if not (controller := hass.data[UNIFI_DOMAIN].get(config_entry.entry_id)): if not (hub := hass.data[UNIFI_DOMAIN].get(config_entry.entry_id)):
return return
controller.load_config_entry_options() hub.load_config_entry_options()
async_dispatcher_send(hass, controller.signal_options_update) async_dispatcher_send(hass, hub.signal_options_update)
@callback @callback
def start_websocket(self) -> None: def start_websocket(self) -> None:
@ -449,7 +449,7 @@ class UniFiController:
self.ws_task.cancel() self.ws_task.cancel()
async def async_reset(self) -> bool: async def async_reset(self) -> bool:
"""Reset this controller to default state. """Reset this hub to default state.
Will cancel any scheduled setup retry and will unload Will cancel any scheduled setup retry and will unload
the config entry. the config entry.
@ -489,11 +489,11 @@ class UniFiController:
return True return True
async def get_unifi_controller( async def get_unifi_api(
hass: HomeAssistant, hass: HomeAssistant,
config: MappingProxyType[str, Any], config: MappingProxyType[str, Any],
) -> aiounifi.Controller: ) -> aiounifi.Controller:
"""Create a controller object and verify authentication.""" """Create a aiounifi object and verify authentication."""
ssl_context: ssl.SSLContext | Literal[False] = False ssl_context: ssl.SSLContext | Literal[False] = False
if verify_ssl := config.get(CONF_VERIFY_SSL): if verify_ssl := config.get(CONF_VERIFY_SSL):
@ -505,7 +505,7 @@ async def get_unifi_controller(
hass, verify_ssl=False, cookie_jar=CookieJar(unsafe=True) hass, verify_ssl=False, cookie_jar=CookieJar(unsafe=True)
) )
controller = aiounifi.Controller( api = aiounifi.Controller(
Configuration( Configuration(
session, session,
host=config[CONF_HOST], host=config[CONF_HOST],
@ -519,8 +519,8 @@ async def get_unifi_controller(
try: try:
async with asyncio.timeout(10): async with asyncio.timeout(10):
await controller.login() await api.login()
return controller return api
except aiounifi.Unauthorized as err: except aiounifi.Unauthorized as err:
LOGGER.warning( LOGGER.warning(

View file

@ -20,7 +20,6 @@ from homeassistant.core import HomeAssistant, callback
from homeassistant.helpers.entity_platform import AddEntitiesCallback from homeassistant.helpers.entity_platform import AddEntitiesCallback
import homeassistant.util.dt as dt_util import homeassistant.util.dt as dt_util
from .controller import UniFiController
from .entity import ( from .entity import (
HandlerT, HandlerT,
UnifiEntity, UnifiEntity,
@ -28,19 +27,20 @@ from .entity import (
async_wlan_available_fn, async_wlan_available_fn,
async_wlan_device_info_fn, async_wlan_device_info_fn,
) )
from .hub import UnifiHub
@callback @callback
def async_wlan_qr_code_image_fn(controller: UniFiController, wlan: Wlan) -> bytes: def async_wlan_qr_code_image_fn(hub: UnifiHub, wlan: Wlan) -> bytes:
"""Calculate receiving data transfer value.""" """Calculate receiving data transfer value."""
return controller.api.wlans.generate_wlan_qr_code(wlan) return hub.api.wlans.generate_wlan_qr_code(wlan)
@dataclass(frozen=True) @dataclass(frozen=True)
class UnifiImageEntityDescriptionMixin(Generic[HandlerT, ApiItemT]): class UnifiImageEntityDescriptionMixin(Generic[HandlerT, ApiItemT]):
"""Validate and load entities from different UniFi handlers.""" """Validate and load entities from different UniFi handlers."""
image_fn: Callable[[UniFiController, ApiItemT], bytes] image_fn: Callable[[UnifiHub, ApiItemT], bytes]
value_fn: Callable[[ApiItemT], str | None] value_fn: Callable[[ApiItemT], str | None]
@ -59,7 +59,7 @@ ENTITY_DESCRIPTIONS: tuple[UnifiImageEntityDescription, ...] = (
entity_category=EntityCategory.DIAGNOSTIC, entity_category=EntityCategory.DIAGNOSTIC,
has_entity_name=True, has_entity_name=True,
entity_registry_enabled_default=False, entity_registry_enabled_default=False,
allowed_fn=lambda controller, obj_id: True, allowed_fn=lambda hub, obj_id: True,
api_handler_fn=lambda api: api.wlans, api_handler_fn=lambda api: api.wlans,
available_fn=async_wlan_available_fn, available_fn=async_wlan_available_fn,
device_info_fn=async_wlan_device_info_fn, device_info_fn=async_wlan_device_info_fn,
@ -68,8 +68,8 @@ ENTITY_DESCRIPTIONS: tuple[UnifiImageEntityDescription, ...] = (
name_fn=lambda wlan: "QR Code", name_fn=lambda wlan: "QR Code",
object_fn=lambda api, obj_id: api.wlans[obj_id], object_fn=lambda api, obj_id: api.wlans[obj_id],
should_poll=False, should_poll=False,
supported_fn=lambda controller, obj_id: True, supported_fn=lambda hub, obj_id: True,
unique_id_fn=lambda controller, obj_id: f"qr_code-{obj_id}", unique_id_fn=lambda hub, obj_id: f"qr_code-{obj_id}",
image_fn=async_wlan_qr_code_image_fn, image_fn=async_wlan_qr_code_image_fn,
value_fn=lambda obj: obj.x_passphrase, value_fn=lambda obj: obj.x_passphrase,
), ),
@ -82,7 +82,7 @@ async def async_setup_entry(
async_add_entities: AddEntitiesCallback, async_add_entities: AddEntitiesCallback,
) -> None: ) -> None:
"""Set up image platform for UniFi Network integration.""" """Set up image platform for UniFi Network integration."""
UniFiController.register_platform( UnifiHub.register_platform(
hass, hass,
config_entry, config_entry,
async_add_entities, async_add_entities,
@ -104,26 +104,26 @@ class UnifiImageEntity(UnifiEntity[HandlerT, ApiItemT], ImageEntity):
def __init__( def __init__(
self, self,
obj_id: str, obj_id: str,
controller: UniFiController, hub: UnifiHub,
description: UnifiEntityDescription[HandlerT, ApiItemT], description: UnifiEntityDescription[HandlerT, ApiItemT],
) -> None: ) -> None:
"""Initiatlize UniFi Image entity.""" """Initiatlize UniFi Image entity."""
super().__init__(obj_id, controller, description) super().__init__(obj_id, hub, description)
ImageEntity.__init__(self, controller.hass) ImageEntity.__init__(self, hub.hass)
def image(self) -> bytes | None: def image(self) -> bytes | None:
"""Return bytes of image.""" """Return bytes of image."""
if self.current_image is None: if self.current_image is None:
description = self.entity_description description = self.entity_description
obj = description.object_fn(self.controller.api, self._obj_id) obj = description.object_fn(self.hub.api, self._obj_id)
self.current_image = description.image_fn(self.controller, obj) self.current_image = description.image_fn(self.hub, obj)
return self.current_image return self.current_image
@callback @callback
def async_update_state(self, event: ItemEvent, obj_id: str) -> None: def async_update_state(self, event: ItemEvent, obj_id: str) -> None:
"""Update entity state.""" """Update entity state."""
description = self.entity_description description = self.entity_description
obj = description.object_fn(self.controller.api, self._obj_id) obj = description.object_fn(self.hub.api, self._obj_id)
if (value := description.value_fn(obj)) != self.previous_value: if (value := description.value_fn(obj)) != self.previous_value:
self.previous_value = value self.previous_value = value
self.current_image = None self.current_image = None

View file

@ -40,7 +40,6 @@ from homeassistant.helpers.typing import StateType
import homeassistant.util.dt as dt_util import homeassistant.util.dt as dt_util
from .const import DEVICE_STATES from .const import DEVICE_STATES
from .controller import UniFiController
from .entity import ( from .entity import (
HandlerT, HandlerT,
UnifiEntity, UnifiEntity,
@ -51,44 +50,43 @@ from .entity import (
async_wlan_available_fn, async_wlan_available_fn,
async_wlan_device_info_fn, async_wlan_device_info_fn,
) )
from .hub import UnifiHub
@callback @callback
def async_bandwidth_sensor_allowed_fn(controller: UniFiController, obj_id: str) -> bool: def async_bandwidth_sensor_allowed_fn(hub: UnifiHub, obj_id: str) -> bool:
"""Check if client is allowed.""" """Check if client is allowed."""
if obj_id in controller.option_supported_clients: if obj_id in hub.option_supported_clients:
return True return True
return controller.option_allow_bandwidth_sensors return hub.option_allow_bandwidth_sensors
@callback @callback
def async_uptime_sensor_allowed_fn(controller: UniFiController, obj_id: str) -> bool: def async_uptime_sensor_allowed_fn(hub: UnifiHub, obj_id: str) -> bool:
"""Check if client is allowed.""" """Check if client is allowed."""
if obj_id in controller.option_supported_clients: if obj_id in hub.option_supported_clients:
return True return True
return controller.option_allow_uptime_sensors return hub.option_allow_uptime_sensors
@callback @callback
def async_client_rx_value_fn(controller: UniFiController, client: Client) -> float: def async_client_rx_value_fn(hub: UnifiHub, client: Client) -> float:
"""Calculate receiving data transfer value.""" """Calculate receiving data transfer value."""
if controller.wireless_clients.is_wireless(client): if hub.wireless_clients.is_wireless(client):
return client.rx_bytes_r / 1000000 return client.rx_bytes_r / 1000000
return client.wired_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(hub: UnifiHub, client: Client) -> float:
"""Calculate transmission data transfer value.""" """Calculate transmission data transfer value."""
if controller.wireless_clients.is_wireless(client): if hub.wireless_clients.is_wireless(client):
return client.tx_bytes_r / 1000000 return client.tx_bytes_r / 1000000
return client.wired_tx_bytes_r / 1000000 return client.wired_tx_bytes_r / 1000000
@callback @callback
def async_client_uptime_value_fn( def async_client_uptime_value_fn(hub: UnifiHub, client: Client) -> datetime:
controller: UniFiController, client: Client
) -> datetime:
"""Calculate the uptime of the client.""" """Calculate the uptime of the client."""
if client.uptime < 1000000000: if client.uptime < 1000000000:
return dt_util.now() - timedelta(seconds=client.uptime) return dt_util.now() - timedelta(seconds=client.uptime)
@ -96,23 +94,21 @@ def async_client_uptime_value_fn(
@callback @callback
def async_wlan_client_value_fn(controller: UniFiController, wlan: Wlan) -> int: def async_wlan_client_value_fn(hub: UnifiHub, wlan: Wlan) -> int:
"""Calculate the amount of clients connected to a wlan.""" """Calculate the amount of clients connected to a wlan."""
return len( return len(
[ [
client.mac client.mac
for client in controller.api.clients.values() for client in hub.api.clients.values()
if client.essid == wlan.name if client.essid == wlan.name
and dt_util.utcnow() - dt_util.utc_from_timestamp(client.last_seen or 0) and dt_util.utcnow() - dt_util.utc_from_timestamp(client.last_seen or 0)
< controller.option_detection_time < hub.option_detection_time
] ]
) )
@callback @callback
def async_device_uptime_value_fn( def async_device_uptime_value_fn(hub: UnifiHub, device: Device) -> datetime | None:
controller: UniFiController, device: Device
) -> datetime | None:
"""Calculate the approximate time the device started (based on uptime returned from API, in seconds).""" """Calculate the approximate time the device started (based on uptime returned from API, in seconds)."""
if device.uptime <= 0: if device.uptime <= 0:
# Library defaults to 0 if uptime is not provided, e.g. when offline # Library defaults to 0 if uptime is not provided, e.g. when offline
@ -131,29 +127,27 @@ def async_device_uptime_value_changed_fn(
@callback @callback
def async_device_outlet_power_supported_fn( def async_device_outlet_power_supported_fn(hub: UnifiHub, obj_id: str) -> bool:
controller: UniFiController, obj_id: str
) -> bool:
"""Determine if an outlet has the power property.""" """Determine if an outlet has the power property."""
# At this time, an outlet_caps value of 3 is expected to indicate that the outlet # At this time, an outlet_caps value of 3 is expected to indicate that the outlet
# supports metering # supports metering
return controller.api.outlets[obj_id].caps == 3 return hub.api.outlets[obj_id].caps == 3
@callback @callback
def async_device_outlet_supported_fn(controller: UniFiController, obj_id: str) -> bool: def async_device_outlet_supported_fn(hub: UnifiHub, obj_id: str) -> bool:
"""Determine if a device supports reading overall power metrics.""" """Determine if a device supports reading overall power metrics."""
return controller.api.devices[obj_id].outlet_ac_power_budget is not None return hub.api.devices[obj_id].outlet_ac_power_budget is not None
@callback @callback
def async_client_is_connected_fn(controller: UniFiController, obj_id: str) -> bool: def async_client_is_connected_fn(hub: UnifiHub, obj_id: str) -> bool:
"""Check if client was last seen recently.""" """Check if client was last seen recently."""
client = controller.api.clients[obj_id] client = hub.api.clients[obj_id]
if ( if (
dt_util.utcnow() - dt_util.utc_from_timestamp(client.last_seen or 0) dt_util.utcnow() - dt_util.utc_from_timestamp(client.last_seen or 0)
> controller.option_detection_time > hub.option_detection_time
): ):
return False return False
@ -164,11 +158,11 @@ def async_client_is_connected_fn(controller: UniFiController, obj_id: str) -> bo
class UnifiSensorEntityDescriptionMixin(Generic[HandlerT, ApiItemT]): class UnifiSensorEntityDescriptionMixin(Generic[HandlerT, ApiItemT]):
"""Validate and load entities from different UniFi handlers.""" """Validate and load entities from different UniFi handlers."""
value_fn: Callable[[UniFiController, ApiItemT], datetime | float | str | None] value_fn: Callable[[UnifiHub, ApiItemT], datetime | float | str | None]
@callback @callback
def async_device_state_value_fn(controller: UniFiController, device: Device) -> str: def async_device_state_value_fn(hub: UnifiHub, device: Device) -> str:
"""Retrieve the state of the device.""" """Retrieve the state of the device."""
return DEVICE_STATES[device.state] return DEVICE_STATES[device.state]
@ -181,7 +175,7 @@ class UnifiSensorEntityDescription(
): ):
"""Class describing UniFi sensor entity.""" """Class describing UniFi sensor entity."""
is_connected_fn: Callable[[UniFiController, str], bool] | None = None is_connected_fn: Callable[[UnifiHub, str], bool] | None = None
# Custom function to determine whether a state change should be recorded # Custom function to determine whether a state change should be recorded
value_changed_fn: Callable[ value_changed_fn: Callable[
[StateType | date | datetime | Decimal, datetime | float | str | None], [StateType | date | datetime | Decimal, datetime | float | str | None],
@ -200,7 +194,7 @@ ENTITY_DESCRIPTIONS: tuple[UnifiSensorEntityDescription, ...] = (
has_entity_name=True, has_entity_name=True,
allowed_fn=async_bandwidth_sensor_allowed_fn, allowed_fn=async_bandwidth_sensor_allowed_fn,
api_handler_fn=lambda api: api.clients, api_handler_fn=lambda api: api.clients,
available_fn=lambda controller, _: controller.available, available_fn=lambda hub, _: hub.available,
device_info_fn=async_client_device_info_fn, device_info_fn=async_client_device_info_fn,
event_is_on=None, event_is_on=None,
event_to_subscribe=None, event_to_subscribe=None,
@ -208,8 +202,8 @@ ENTITY_DESCRIPTIONS: tuple[UnifiSensorEntityDescription, ...] = (
name_fn=lambda _: "RX", name_fn=lambda _: "RX",
object_fn=lambda api, obj_id: api.clients[obj_id], object_fn=lambda api, obj_id: api.clients[obj_id],
should_poll=False, should_poll=False,
supported_fn=lambda controller, _: controller.option_allow_bandwidth_sensors, supported_fn=lambda hub, _: hub.option_allow_bandwidth_sensors,
unique_id_fn=lambda controller, obj_id: f"rx-{obj_id}", unique_id_fn=lambda hub, obj_id: f"rx-{obj_id}",
value_fn=async_client_rx_value_fn, value_fn=async_client_rx_value_fn,
), ),
UnifiSensorEntityDescription[Clients, Client]( UnifiSensorEntityDescription[Clients, Client](
@ -222,7 +216,7 @@ ENTITY_DESCRIPTIONS: tuple[UnifiSensorEntityDescription, ...] = (
has_entity_name=True, has_entity_name=True,
allowed_fn=async_bandwidth_sensor_allowed_fn, allowed_fn=async_bandwidth_sensor_allowed_fn,
api_handler_fn=lambda api: api.clients, api_handler_fn=lambda api: api.clients,
available_fn=lambda controller, _: controller.available, available_fn=lambda hub, _: hub.available,
device_info_fn=async_client_device_info_fn, device_info_fn=async_client_device_info_fn,
event_is_on=None, event_is_on=None,
event_to_subscribe=None, event_to_subscribe=None,
@ -230,8 +224,8 @@ ENTITY_DESCRIPTIONS: tuple[UnifiSensorEntityDescription, ...] = (
name_fn=lambda _: "TX", name_fn=lambda _: "TX",
object_fn=lambda api, obj_id: api.clients[obj_id], object_fn=lambda api, obj_id: api.clients[obj_id],
should_poll=False, should_poll=False,
supported_fn=lambda controller, _: controller.option_allow_bandwidth_sensors, supported_fn=lambda hub, _: hub.option_allow_bandwidth_sensors,
unique_id_fn=lambda controller, obj_id: f"tx-{obj_id}", unique_id_fn=lambda hub, obj_id: f"tx-{obj_id}",
value_fn=async_client_tx_value_fn, value_fn=async_client_tx_value_fn,
), ),
UnifiSensorEntityDescription[Ports, Port]( UnifiSensorEntityDescription[Ports, Port](
@ -241,7 +235,7 @@ ENTITY_DESCRIPTIONS: tuple[UnifiSensorEntityDescription, ...] = (
native_unit_of_measurement=UnitOfPower.WATT, native_unit_of_measurement=UnitOfPower.WATT,
has_entity_name=True, has_entity_name=True,
entity_registry_enabled_default=False, entity_registry_enabled_default=False,
allowed_fn=lambda controller, obj_id: True, allowed_fn=lambda hub, obj_id: True,
api_handler_fn=lambda api: api.ports, api_handler_fn=lambda api: api.ports,
available_fn=async_device_available_fn, available_fn=async_device_available_fn,
device_info_fn=async_device_device_info_fn, device_info_fn=async_device_device_info_fn,
@ -250,8 +244,8 @@ ENTITY_DESCRIPTIONS: tuple[UnifiSensorEntityDescription, ...] = (
name_fn=lambda port: f"{port.name} PoE Power", name_fn=lambda port: f"{port.name} PoE Power",
object_fn=lambda api, obj_id: api.ports[obj_id], object_fn=lambda api, obj_id: api.ports[obj_id],
should_poll=False, should_poll=False,
supported_fn=lambda controller, obj_id: controller.api.ports[obj_id].port_poe, supported_fn=lambda hub, obj_id: hub.api.ports[obj_id].port_poe,
unique_id_fn=lambda controller, obj_id: f"poe_power-{obj_id}", unique_id_fn=lambda hub, obj_id: f"poe_power-{obj_id}",
value_fn=lambda _, obj: obj.poe_power if obj.poe_mode != "off" else "0", value_fn=lambda _, obj: obj.poe_power if obj.poe_mode != "off" else "0",
), ),
UnifiSensorEntityDescription[Clients, Client]( UnifiSensorEntityDescription[Clients, Client](
@ -262,15 +256,15 @@ ENTITY_DESCRIPTIONS: tuple[UnifiSensorEntityDescription, ...] = (
entity_registry_enabled_default=False, entity_registry_enabled_default=False,
allowed_fn=async_uptime_sensor_allowed_fn, allowed_fn=async_uptime_sensor_allowed_fn,
api_handler_fn=lambda api: api.clients, api_handler_fn=lambda api: api.clients,
available_fn=lambda controller, obj_id: controller.available, available_fn=lambda hub, obj_id: hub.available,
device_info_fn=async_client_device_info_fn, device_info_fn=async_client_device_info_fn,
event_is_on=None, event_is_on=None,
event_to_subscribe=None, event_to_subscribe=None,
name_fn=lambda client: "Uptime", name_fn=lambda client: "Uptime",
object_fn=lambda api, obj_id: api.clients[obj_id], object_fn=lambda api, obj_id: api.clients[obj_id],
should_poll=False, should_poll=False,
supported_fn=lambda controller, _: controller.option_allow_uptime_sensors, supported_fn=lambda hub, _: hub.option_allow_uptime_sensors,
unique_id_fn=lambda controller, obj_id: f"uptime-{obj_id}", unique_id_fn=lambda hub, obj_id: f"uptime-{obj_id}",
value_fn=async_client_uptime_value_fn, value_fn=async_client_uptime_value_fn,
), ),
UnifiSensorEntityDescription[Wlans, Wlan]( UnifiSensorEntityDescription[Wlans, Wlan](
@ -278,7 +272,7 @@ ENTITY_DESCRIPTIONS: tuple[UnifiSensorEntityDescription, ...] = (
entity_category=EntityCategory.DIAGNOSTIC, entity_category=EntityCategory.DIAGNOSTIC,
has_entity_name=True, has_entity_name=True,
state_class=SensorStateClass.MEASUREMENT, state_class=SensorStateClass.MEASUREMENT,
allowed_fn=lambda controller, obj_id: True, allowed_fn=lambda hub, obj_id: True,
api_handler_fn=lambda api: api.wlans, api_handler_fn=lambda api: api.wlans,
available_fn=async_wlan_available_fn, available_fn=async_wlan_available_fn,
device_info_fn=async_wlan_device_info_fn, device_info_fn=async_wlan_device_info_fn,
@ -287,8 +281,8 @@ ENTITY_DESCRIPTIONS: tuple[UnifiSensorEntityDescription, ...] = (
name_fn=lambda wlan: None, name_fn=lambda wlan: None,
object_fn=lambda api, obj_id: api.wlans[obj_id], object_fn=lambda api, obj_id: api.wlans[obj_id],
should_poll=True, should_poll=True,
supported_fn=lambda controller, obj_id: True, supported_fn=lambda hub, obj_id: True,
unique_id_fn=lambda controller, obj_id: f"wlan_clients-{obj_id}", unique_id_fn=lambda hub, obj_id: f"wlan_clients-{obj_id}",
value_fn=async_wlan_client_value_fn, value_fn=async_wlan_client_value_fn,
), ),
UnifiSensorEntityDescription[Outlets, Outlet]( UnifiSensorEntityDescription[Outlets, Outlet](
@ -297,7 +291,7 @@ ENTITY_DESCRIPTIONS: tuple[UnifiSensorEntityDescription, ...] = (
entity_category=EntityCategory.DIAGNOSTIC, entity_category=EntityCategory.DIAGNOSTIC,
native_unit_of_measurement=UnitOfPower.WATT, native_unit_of_measurement=UnitOfPower.WATT,
has_entity_name=True, has_entity_name=True,
allowed_fn=lambda controller, obj_id: True, allowed_fn=lambda hub, obj_id: True,
api_handler_fn=lambda api: api.outlets, api_handler_fn=lambda api: api.outlets,
available_fn=async_device_available_fn, available_fn=async_device_available_fn,
device_info_fn=async_device_device_info_fn, device_info_fn=async_device_device_info_fn,
@ -307,7 +301,7 @@ ENTITY_DESCRIPTIONS: tuple[UnifiSensorEntityDescription, ...] = (
object_fn=lambda api, obj_id: api.outlets[obj_id], object_fn=lambda api, obj_id: api.outlets[obj_id],
should_poll=True, should_poll=True,
supported_fn=async_device_outlet_power_supported_fn, supported_fn=async_device_outlet_power_supported_fn,
unique_id_fn=lambda controller, obj_id: f"outlet_power-{obj_id}", unique_id_fn=lambda hub, obj_id: f"outlet_power-{obj_id}",
value_fn=lambda _, obj: obj.power if obj.relay_state else "0", value_fn=lambda _, obj: obj.power if obj.relay_state else "0",
), ),
UnifiSensorEntityDescription[Devices, Device]( UnifiSensorEntityDescription[Devices, Device](
@ -317,7 +311,7 @@ ENTITY_DESCRIPTIONS: tuple[UnifiSensorEntityDescription, ...] = (
native_unit_of_measurement=UnitOfPower.WATT, native_unit_of_measurement=UnitOfPower.WATT,
suggested_display_precision=1, suggested_display_precision=1,
has_entity_name=True, has_entity_name=True,
allowed_fn=lambda controller, obj_id: True, allowed_fn=lambda hub, obj_id: True,
api_handler_fn=lambda api: api.devices, api_handler_fn=lambda api: api.devices,
available_fn=async_device_available_fn, available_fn=async_device_available_fn,
device_info_fn=async_device_device_info_fn, device_info_fn=async_device_device_info_fn,
@ -327,8 +321,8 @@ ENTITY_DESCRIPTIONS: tuple[UnifiSensorEntityDescription, ...] = (
object_fn=lambda api, obj_id: api.devices[obj_id], object_fn=lambda api, obj_id: api.devices[obj_id],
should_poll=False, should_poll=False,
supported_fn=async_device_outlet_supported_fn, supported_fn=async_device_outlet_supported_fn,
unique_id_fn=lambda controller, obj_id: f"ac_power_budget-{obj_id}", unique_id_fn=lambda hub, obj_id: f"ac_power_budget-{obj_id}",
value_fn=lambda controller, device: device.outlet_ac_power_budget, value_fn=lambda hub, device: device.outlet_ac_power_budget,
), ),
UnifiSensorEntityDescription[Devices, Device]( UnifiSensorEntityDescription[Devices, Device](
key="SmartPower AC power consumption", key="SmartPower AC power consumption",
@ -337,7 +331,7 @@ ENTITY_DESCRIPTIONS: tuple[UnifiSensorEntityDescription, ...] = (
native_unit_of_measurement=UnitOfPower.WATT, native_unit_of_measurement=UnitOfPower.WATT,
suggested_display_precision=1, suggested_display_precision=1,
has_entity_name=True, has_entity_name=True,
allowed_fn=lambda controller, obj_id: True, allowed_fn=lambda hub, obj_id: True,
api_handler_fn=lambda api: api.devices, api_handler_fn=lambda api: api.devices,
available_fn=async_device_available_fn, available_fn=async_device_available_fn,
device_info_fn=async_device_device_info_fn, device_info_fn=async_device_device_info_fn,
@ -347,15 +341,15 @@ ENTITY_DESCRIPTIONS: tuple[UnifiSensorEntityDescription, ...] = (
object_fn=lambda api, obj_id: api.devices[obj_id], object_fn=lambda api, obj_id: api.devices[obj_id],
should_poll=False, should_poll=False,
supported_fn=async_device_outlet_supported_fn, supported_fn=async_device_outlet_supported_fn,
unique_id_fn=lambda controller, obj_id: f"ac_power_conumption-{obj_id}", unique_id_fn=lambda hub, obj_id: f"ac_power_conumption-{obj_id}",
value_fn=lambda controller, device: device.outlet_ac_power_consumption, value_fn=lambda hub, device: device.outlet_ac_power_consumption,
), ),
UnifiSensorEntityDescription[Devices, Device]( UnifiSensorEntityDescription[Devices, Device](
key="Device uptime", key="Device uptime",
device_class=SensorDeviceClass.TIMESTAMP, device_class=SensorDeviceClass.TIMESTAMP,
entity_category=EntityCategory.DIAGNOSTIC, entity_category=EntityCategory.DIAGNOSTIC,
has_entity_name=True, has_entity_name=True,
allowed_fn=lambda controller, obj_id: True, allowed_fn=lambda hub, obj_id: True,
api_handler_fn=lambda api: api.devices, api_handler_fn=lambda api: api.devices,
available_fn=async_device_available_fn, available_fn=async_device_available_fn,
device_info_fn=async_device_device_info_fn, device_info_fn=async_device_device_info_fn,
@ -364,8 +358,8 @@ ENTITY_DESCRIPTIONS: tuple[UnifiSensorEntityDescription, ...] = (
name_fn=lambda device: "Uptime", name_fn=lambda device: "Uptime",
object_fn=lambda api, obj_id: api.devices[obj_id], object_fn=lambda api, obj_id: api.devices[obj_id],
should_poll=False, should_poll=False,
supported_fn=lambda controller, obj_id: True, supported_fn=lambda hub, obj_id: True,
unique_id_fn=lambda controller, obj_id: f"device_uptime-{obj_id}", unique_id_fn=lambda hub, obj_id: f"device_uptime-{obj_id}",
value_fn=async_device_uptime_value_fn, value_fn=async_device_uptime_value_fn,
value_changed_fn=async_device_uptime_value_changed_fn, value_changed_fn=async_device_uptime_value_changed_fn,
), ),
@ -375,7 +369,7 @@ ENTITY_DESCRIPTIONS: tuple[UnifiSensorEntityDescription, ...] = (
entity_category=EntityCategory.DIAGNOSTIC, entity_category=EntityCategory.DIAGNOSTIC,
native_unit_of_measurement=UnitOfTemperature.CELSIUS, native_unit_of_measurement=UnitOfTemperature.CELSIUS,
has_entity_name=True, has_entity_name=True,
allowed_fn=lambda controller, obj_id: True, allowed_fn=lambda hub, obj_id: True,
api_handler_fn=lambda api: api.devices, api_handler_fn=lambda api: api.devices,
available_fn=async_device_available_fn, available_fn=async_device_available_fn,
device_info_fn=async_device_device_info_fn, device_info_fn=async_device_device_info_fn,
@ -385,7 +379,7 @@ ENTITY_DESCRIPTIONS: tuple[UnifiSensorEntityDescription, ...] = (
object_fn=lambda api, obj_id: api.devices[obj_id], object_fn=lambda api, obj_id: api.devices[obj_id],
should_poll=False, should_poll=False,
supported_fn=lambda ctrlr, obj_id: ctrlr.api.devices[obj_id].has_temperature, supported_fn=lambda ctrlr, obj_id: ctrlr.api.devices[obj_id].has_temperature,
unique_id_fn=lambda controller, obj_id: f"device_temperature-{obj_id}", unique_id_fn=lambda hub, obj_id: f"device_temperature-{obj_id}",
value_fn=lambda ctrlr, device: device.general_temperature, value_fn=lambda ctrlr, device: device.general_temperature,
), ),
UnifiSensorEntityDescription[Devices, Device]( UnifiSensorEntityDescription[Devices, Device](
@ -393,7 +387,7 @@ ENTITY_DESCRIPTIONS: tuple[UnifiSensorEntityDescription, ...] = (
device_class=SensorDeviceClass.ENUM, device_class=SensorDeviceClass.ENUM,
entity_category=EntityCategory.DIAGNOSTIC, entity_category=EntityCategory.DIAGNOSTIC,
has_entity_name=True, has_entity_name=True,
allowed_fn=lambda controller, obj_id: True, allowed_fn=lambda hub, obj_id: True,
api_handler_fn=lambda api: api.devices, api_handler_fn=lambda api: api.devices,
available_fn=async_device_available_fn, available_fn=async_device_available_fn,
device_info_fn=async_device_device_info_fn, device_info_fn=async_device_device_info_fn,
@ -402,8 +396,8 @@ ENTITY_DESCRIPTIONS: tuple[UnifiSensorEntityDescription, ...] = (
name_fn=lambda device: "State", name_fn=lambda device: "State",
object_fn=lambda api, obj_id: api.devices[obj_id], object_fn=lambda api, obj_id: api.devices[obj_id],
should_poll=False, should_poll=False,
supported_fn=lambda controller, obj_id: True, supported_fn=lambda hub, obj_id: True,
unique_id_fn=lambda controller, obj_id: f"device_state-{obj_id}", unique_id_fn=lambda hub, obj_id: f"device_state-{obj_id}",
value_fn=async_device_state_value_fn, value_fn=async_device_state_value_fn,
options=list(DEVICE_STATES.values()), options=list(DEVICE_STATES.values()),
), ),
@ -416,7 +410,7 @@ async def async_setup_entry(
async_add_entities: AddEntitiesCallback, async_add_entities: AddEntitiesCallback,
) -> None: ) -> None:
"""Set up sensors for UniFi Network integration.""" """Set up sensors for UniFi Network integration."""
UniFiController.register_platform( UnifiHub.register_platform(
hass, config_entry, async_add_entities, UnifiSensorEntity, ENTITY_DESCRIPTIONS hass, config_entry, async_add_entities, UnifiSensorEntity, ENTITY_DESCRIPTIONS
) )
@ -443,19 +437,19 @@ class UnifiSensorEntity(UnifiEntity[HandlerT, ApiItemT], SensorEntity):
Update native_value. Update native_value.
""" """
description = self.entity_description description = self.entity_description
obj = description.object_fn(self.controller.api, self._obj_id) obj = description.object_fn(self.hub.api, self._obj_id)
# Update the value only if value is considered to have changed relative to its previous state # Update the value only if value is considered to have changed relative to its previous state
if description.value_changed_fn( if description.value_changed_fn(
self.native_value, (value := description.value_fn(self.controller, obj)) self.native_value, (value := description.value_fn(self.hub, obj))
): ):
self._attr_native_value = value self._attr_native_value = value
if description.is_connected_fn is not None: if description.is_connected_fn is not None:
# Send heartbeat if client is connected # Send heartbeat if client is connected
if description.is_connected_fn(self.controller, self._obj_id): if description.is_connected_fn(self.hub, self._obj_id):
self.controller.async_heartbeat( self.hub.async_heartbeat(
self._attr_unique_id, self._attr_unique_id,
dt_util.utcnow() + self.controller.option_detection_time, dt_util.utcnow() + self.hub.option_detection_time,
) )
async def async_added_to_hass(self) -> None: async def async_added_to_hass(self) -> None:
@ -467,7 +461,7 @@ class UnifiSensorEntity(UnifiEntity[HandlerT, ApiItemT], SensorEntity):
self.async_on_remove( self.async_on_remove(
async_dispatcher_connect( async_dispatcher_connect(
self.hass, self.hass,
f"{self.controller.signal_heartbeat_missed}_{self.unique_id}", f"{self.hub.signal_heartbeat_missed}_{self.unique_id}",
self._make_disconnected, self._make_disconnected,
) )
) )
@ -478,4 +472,4 @@ class UnifiSensorEntity(UnifiEntity[HandlerT, ApiItemT], SensorEntity):
if self.entity_description.is_connected_fn is not None: if self.entity_description.is_connected_fn is not None:
# Remove heartbeat registration # Remove heartbeat registration
self.controller.async_heartbeat(self._attr_unique_id) self.hub.async_heartbeat(self._attr_unique_id)

View file

@ -72,31 +72,31 @@ async def async_reconnect_client(hass: HomeAssistant, data: Mapping[str, Any]) -
if mac == "": if mac == "":
return return
for controller in hass.data[UNIFI_DOMAIN].values(): for hub in hass.data[UNIFI_DOMAIN].values():
if ( if (
not controller.available not hub.available
or (client := controller.api.clients.get(mac)) is None or (client := hub.api.clients.get(mac)) is None
or client.is_wired or client.is_wired
): ):
continue continue
await controller.api.request(ClientReconnectRequest.create(mac)) await hub.api.request(ClientReconnectRequest.create(mac))
async def async_remove_clients(hass: HomeAssistant, data: Mapping[str, Any]) -> None: async def async_remove_clients(hass: HomeAssistant, data: Mapping[str, Any]) -> None:
"""Remove select clients from controller. """Remove select clients from UniFi Network.
Validates based on: Validates based on:
- Total time between first seen and last seen is less than 15 minutes. - Total time between first seen and last seen is less than 15 minutes.
- Neither IP, hostname nor name is configured. - Neither IP, hostname nor name is configured.
""" """
for controller in hass.data[UNIFI_DOMAIN].values(): for hub in hass.data[UNIFI_DOMAIN].values():
if not controller.available: if not hub.available:
continue continue
clients_to_remove = [] clients_to_remove = []
for client in controller.api.clients_all.values(): for client in hub.api.clients_all.values():
if ( if (
client.last_seen client.last_seen
and client.first_seen and client.first_seen
@ -110,4 +110,4 @@ async def async_remove_clients(hass: HomeAssistant, data: Mapping[str, Any]) ->
clients_to_remove.append(client.mac) clients_to_remove.append(client.mac)
if clients_to_remove: if clients_to_remove:
await controller.api.request(ClientRemoveRequest.create(clients_to_remove)) await hub.api.request(ClientRemoveRequest.create(clients_to_remove))

View file

@ -45,7 +45,6 @@ from homeassistant.helpers.entity_platform import AddEntitiesCallback
import homeassistant.helpers.entity_registry as er import homeassistant.helpers.entity_registry as er
from .const import ATTR_MANUFACTURER from .const import ATTR_MANUFACTURER
from .controller import UNIFI_DOMAIN, UniFiController
from .entity import ( from .entity import (
HandlerT, HandlerT,
SubscriptionT, SubscriptionT,
@ -56,25 +55,24 @@ from .entity import (
async_device_device_info_fn, async_device_device_info_fn,
async_wlan_device_info_fn, async_wlan_device_info_fn,
) )
from .hub import UNIFI_DOMAIN, UnifiHub
CLIENT_BLOCKED = (EventKey.WIRED_CLIENT_BLOCKED, EventKey.WIRELESS_CLIENT_BLOCKED) CLIENT_BLOCKED = (EventKey.WIRED_CLIENT_BLOCKED, EventKey.WIRELESS_CLIENT_BLOCKED)
CLIENT_UNBLOCKED = (EventKey.WIRED_CLIENT_UNBLOCKED, EventKey.WIRELESS_CLIENT_UNBLOCKED) CLIENT_UNBLOCKED = (EventKey.WIRED_CLIENT_UNBLOCKED, EventKey.WIRELESS_CLIENT_UNBLOCKED)
@callback @callback
def async_block_client_allowed_fn(controller: UniFiController, obj_id: str) -> bool: def async_block_client_allowed_fn(hub: UnifiHub, obj_id: str) -> bool:
"""Check if client is allowed.""" """Check if client is allowed."""
if obj_id in controller.option_supported_clients: if obj_id in hub.option_supported_clients:
return True return True
return obj_id in controller.option_block_clients return obj_id in hub.option_block_clients
@callback @callback
def async_dpi_group_is_on_fn( def async_dpi_group_is_on_fn(hub: UnifiHub, dpi_group: DPIRestrictionGroup) -> bool:
controller: UniFiController, dpi_group: DPIRestrictionGroup
) -> bool:
"""Calculate if all apps are enabled.""" """Calculate if all apps are enabled."""
api = controller.api api = hub.api
return all( return all(
api.dpi_apps[app_id].enabled api.dpi_apps[app_id].enabled
for app_id in dpi_group.dpiapp_ids or [] for app_id in dpi_group.dpiapp_ids or []
@ -83,9 +81,7 @@ def async_dpi_group_is_on_fn(
@callback @callback
def async_dpi_group_device_info_fn( def async_dpi_group_device_info_fn(hub: UnifiHub, obj_id: str) -> DeviceInfo:
controller: UniFiController, obj_id: str
) -> DeviceInfo:
"""Create device registry entry for DPI group.""" """Create device registry entry for DPI group."""
return DeviceInfo( return DeviceInfo(
entry_type=DeviceEntryType.SERVICE, entry_type=DeviceEntryType.SERVICE,
@ -97,11 +93,9 @@ def async_dpi_group_device_info_fn(
@callback @callback
def async_port_forward_device_info_fn( def async_port_forward_device_info_fn(hub: UnifiHub, obj_id: str) -> DeviceInfo:
controller: UniFiController, obj_id: str
) -> DeviceInfo:
"""Create device registry entry for port forward.""" """Create device registry entry for port forward."""
unique_id = controller.config_entry.unique_id unique_id = hub.config_entry.unique_id
assert unique_id is not None assert unique_id is not None
return DeviceInfo( return DeviceInfo(
entry_type=DeviceEntryType.SERVICE, entry_type=DeviceEntryType.SERVICE,
@ -113,79 +107,67 @@ def async_port_forward_device_info_fn(
async def async_block_client_control_fn( async def async_block_client_control_fn(
controller: UniFiController, obj_id: str, target: bool hub: UnifiHub, obj_id: str, target: bool
) -> None: ) -> None:
"""Control network access of client.""" """Control network access of client."""
await controller.api.request(ClientBlockRequest.create(obj_id, not target)) await hub.api.request(ClientBlockRequest.create(obj_id, not target))
async def async_dpi_group_control_fn( async def async_dpi_group_control_fn(hub: UnifiHub, obj_id: str, target: bool) -> None:
controller: UniFiController, obj_id: str, target: bool
) -> None:
"""Enable or disable DPI group.""" """Enable or disable DPI group."""
dpi_group = controller.api.dpi_groups[obj_id] dpi_group = hub.api.dpi_groups[obj_id]
await asyncio.gather( await asyncio.gather(
*[ *[
controller.api.request( hub.api.request(DPIRestrictionAppEnableRequest.create(app_id, target))
DPIRestrictionAppEnableRequest.create(app_id, target)
)
for app_id in dpi_group.dpiapp_ids or [] for app_id in dpi_group.dpiapp_ids or []
] ]
) )
@callback @callback
def async_outlet_supports_switching_fn( def async_outlet_supports_switching_fn(hub: UnifiHub, obj_id: str) -> bool:
controller: UniFiController, obj_id: str
) -> bool:
"""Determine if an outlet supports switching.""" """Determine if an outlet supports switching."""
outlet = controller.api.outlets[obj_id] outlet = hub.api.outlets[obj_id]
return outlet.has_relay or outlet.caps in (1, 3) return outlet.has_relay or outlet.caps in (1, 3)
async def async_outlet_control_fn( async def async_outlet_control_fn(hub: UnifiHub, obj_id: str, target: bool) -> None:
controller: UniFiController, obj_id: str, target: bool
) -> None:
"""Control outlet relay.""" """Control outlet relay."""
mac, _, index = obj_id.partition("_") mac, _, index = obj_id.partition("_")
device = controller.api.devices[mac] device = hub.api.devices[mac]
await controller.api.request( await hub.api.request(
DeviceSetOutletRelayRequest.create(device, int(index), target) DeviceSetOutletRelayRequest.create(device, int(index), target)
) )
async def async_poe_port_control_fn( async def async_poe_port_control_fn(hub: UnifiHub, obj_id: str, target: bool) -> None:
controller: UniFiController, obj_id: str, target: bool
) -> None:
"""Control poe state.""" """Control poe state."""
mac, _, index = obj_id.partition("_") mac, _, index = obj_id.partition("_")
port = controller.api.ports[obj_id] port = hub.api.ports[obj_id]
on_state = "auto" if port.raw["poe_caps"] != 8 else "passthrough" on_state = "auto" if port.raw["poe_caps"] != 8 else "passthrough"
state = on_state if target else "off" state = on_state if target else "off"
controller.async_queue_poe_port_command(mac, int(index), state) hub.async_queue_poe_port_command(mac, int(index), state)
async def async_port_forward_control_fn( async def async_port_forward_control_fn(
controller: UniFiController, obj_id: str, target: bool hub: UnifiHub, obj_id: str, target: bool
) -> None: ) -> None:
"""Control port forward state.""" """Control port forward state."""
port_forward = controller.api.port_forwarding[obj_id] port_forward = hub.api.port_forwarding[obj_id]
await controller.api.request(PortForwardEnableRequest.create(port_forward, target)) await hub.api.request(PortForwardEnableRequest.create(port_forward, target))
async def async_wlan_control_fn( async def async_wlan_control_fn(hub: UnifiHub, obj_id: str, target: bool) -> None:
controller: UniFiController, obj_id: str, target: bool
) -> None:
"""Control outlet relay.""" """Control outlet relay."""
await controller.api.request(WlanEnableRequest.create(obj_id, target)) await hub.api.request(WlanEnableRequest.create(obj_id, target))
@dataclass(frozen=True) @dataclass(frozen=True)
class UnifiSwitchEntityDescriptionMixin(Generic[HandlerT, ApiItemT]): class UnifiSwitchEntityDescriptionMixin(Generic[HandlerT, ApiItemT]):
"""Validate and load entities from different UniFi handlers.""" """Validate and load entities from different UniFi handlers."""
control_fn: Callable[[UniFiController, str, bool], Coroutine[Any, Any, None]] control_fn: Callable[[UnifiHub, str, bool], Coroutine[Any, Any, None]]
is_on_fn: Callable[[UniFiController, ApiItemT], bool] is_on_fn: Callable[[UnifiHub, ApiItemT], bool]
@dataclass(frozen=True) @dataclass(frozen=True)
@ -209,26 +191,26 @@ ENTITY_DESCRIPTIONS: tuple[UnifiSwitchEntityDescription, ...] = (
icon="mdi:ethernet", icon="mdi:ethernet",
allowed_fn=async_block_client_allowed_fn, allowed_fn=async_block_client_allowed_fn,
api_handler_fn=lambda api: api.clients, api_handler_fn=lambda api: api.clients,
available_fn=lambda controller, obj_id: controller.available, available_fn=lambda hub, obj_id: hub.available,
control_fn=async_block_client_control_fn, control_fn=async_block_client_control_fn,
device_info_fn=async_client_device_info_fn, device_info_fn=async_client_device_info_fn,
event_is_on=CLIENT_UNBLOCKED, event_is_on=CLIENT_UNBLOCKED,
event_to_subscribe=CLIENT_BLOCKED + CLIENT_UNBLOCKED, event_to_subscribe=CLIENT_BLOCKED + CLIENT_UNBLOCKED,
is_on_fn=lambda controller, client: not client.blocked, is_on_fn=lambda hub, client: not client.blocked,
name_fn=lambda client: None, name_fn=lambda client: None,
object_fn=lambda api, obj_id: api.clients[obj_id], object_fn=lambda api, obj_id: api.clients[obj_id],
only_event_for_state_change=True, only_event_for_state_change=True,
should_poll=False, should_poll=False,
supported_fn=lambda controller, obj_id: True, supported_fn=lambda hub, obj_id: True,
unique_id_fn=lambda controller, obj_id: f"block-{obj_id}", unique_id_fn=lambda hub, obj_id: f"block-{obj_id}",
), ),
UnifiSwitchEntityDescription[DPIRestrictionGroups, DPIRestrictionGroup]( UnifiSwitchEntityDescription[DPIRestrictionGroups, DPIRestrictionGroup](
key="DPI restriction", key="DPI restriction",
entity_category=EntityCategory.CONFIG, entity_category=EntityCategory.CONFIG,
icon="mdi:network", icon="mdi:network",
allowed_fn=lambda controller, obj_id: controller.option_dpi_restrictions, allowed_fn=lambda hub, obj_id: hub.option_dpi_restrictions,
api_handler_fn=lambda api: api.dpi_groups, api_handler_fn=lambda api: api.dpi_groups,
available_fn=lambda controller, obj_id: controller.available, available_fn=lambda hub, obj_id: hub.available,
control_fn=async_dpi_group_control_fn, control_fn=async_dpi_group_control_fn,
custom_subscribe=lambda api: api.dpi_apps.subscribe, custom_subscribe=lambda api: api.dpi_apps.subscribe,
device_info_fn=async_dpi_group_device_info_fn, device_info_fn=async_dpi_group_device_info_fn,
@ -239,25 +221,25 @@ ENTITY_DESCRIPTIONS: tuple[UnifiSwitchEntityDescription, ...] = (
object_fn=lambda api, obj_id: api.dpi_groups[obj_id], object_fn=lambda api, obj_id: api.dpi_groups[obj_id],
should_poll=False, should_poll=False,
supported_fn=lambda c, obj_id: bool(c.api.dpi_groups[obj_id].dpiapp_ids), supported_fn=lambda c, obj_id: bool(c.api.dpi_groups[obj_id].dpiapp_ids),
unique_id_fn=lambda controller, obj_id: obj_id, unique_id_fn=lambda hub, obj_id: obj_id,
), ),
UnifiSwitchEntityDescription[Outlets, Outlet]( UnifiSwitchEntityDescription[Outlets, Outlet](
key="Outlet control", key="Outlet control",
device_class=SwitchDeviceClass.OUTLET, device_class=SwitchDeviceClass.OUTLET,
has_entity_name=True, has_entity_name=True,
allowed_fn=lambda controller, obj_id: True, allowed_fn=lambda hub, obj_id: True,
api_handler_fn=lambda api: api.outlets, api_handler_fn=lambda api: api.outlets,
available_fn=async_device_available_fn, available_fn=async_device_available_fn,
control_fn=async_outlet_control_fn, control_fn=async_outlet_control_fn,
device_info_fn=async_device_device_info_fn, device_info_fn=async_device_device_info_fn,
event_is_on=None, event_is_on=None,
event_to_subscribe=None, event_to_subscribe=None,
is_on_fn=lambda controller, outlet: outlet.relay_state, is_on_fn=lambda hub, outlet: outlet.relay_state,
name_fn=lambda outlet: outlet.name, name_fn=lambda outlet: outlet.name,
object_fn=lambda api, obj_id: api.outlets[obj_id], object_fn=lambda api, obj_id: api.outlets[obj_id],
should_poll=False, should_poll=False,
supported_fn=async_outlet_supports_switching_fn, supported_fn=async_outlet_supports_switching_fn,
unique_id_fn=lambda controller, obj_id: f"outlet-{obj_id}", unique_id_fn=lambda hub, obj_id: f"outlet-{obj_id}",
), ),
UnifiSwitchEntityDescription[PortForwarding, PortForward]( UnifiSwitchEntityDescription[PortForwarding, PortForward](
key="Port forward control", key="Port forward control",
@ -265,19 +247,19 @@ ENTITY_DESCRIPTIONS: tuple[UnifiSwitchEntityDescription, ...] = (
entity_category=EntityCategory.CONFIG, entity_category=EntityCategory.CONFIG,
has_entity_name=True, has_entity_name=True,
icon="mdi:upload-network", icon="mdi:upload-network",
allowed_fn=lambda controller, obj_id: True, allowed_fn=lambda hub, obj_id: True,
api_handler_fn=lambda api: api.port_forwarding, api_handler_fn=lambda api: api.port_forwarding,
available_fn=lambda controller, obj_id: controller.available, available_fn=lambda hub, obj_id: hub.available,
control_fn=async_port_forward_control_fn, control_fn=async_port_forward_control_fn,
device_info_fn=async_port_forward_device_info_fn, device_info_fn=async_port_forward_device_info_fn,
event_is_on=None, event_is_on=None,
event_to_subscribe=None, event_to_subscribe=None,
is_on_fn=lambda controller, port_forward: port_forward.enabled, is_on_fn=lambda hub, port_forward: port_forward.enabled,
name_fn=lambda port_forward: f"{port_forward.name}", name_fn=lambda port_forward: f"{port_forward.name}",
object_fn=lambda api, obj_id: api.port_forwarding[obj_id], object_fn=lambda api, obj_id: api.port_forwarding[obj_id],
should_poll=False, should_poll=False,
supported_fn=lambda controller, obj_id: True, supported_fn=lambda hub, obj_id: True,
unique_id_fn=lambda controller, obj_id: f"port_forward-{obj_id}", unique_id_fn=lambda hub, obj_id: f"port_forward-{obj_id}",
), ),
UnifiSwitchEntityDescription[Ports, Port]( UnifiSwitchEntityDescription[Ports, Port](
key="PoE port control", key="PoE port control",
@ -286,19 +268,19 @@ ENTITY_DESCRIPTIONS: tuple[UnifiSwitchEntityDescription, ...] = (
has_entity_name=True, has_entity_name=True,
entity_registry_enabled_default=False, entity_registry_enabled_default=False,
icon="mdi:ethernet", icon="mdi:ethernet",
allowed_fn=lambda controller, obj_id: True, allowed_fn=lambda hub, obj_id: True,
api_handler_fn=lambda api: api.ports, api_handler_fn=lambda api: api.ports,
available_fn=async_device_available_fn, available_fn=async_device_available_fn,
control_fn=async_poe_port_control_fn, control_fn=async_poe_port_control_fn,
device_info_fn=async_device_device_info_fn, device_info_fn=async_device_device_info_fn,
event_is_on=None, event_is_on=None,
event_to_subscribe=None, event_to_subscribe=None,
is_on_fn=lambda controller, port: port.poe_mode != "off", is_on_fn=lambda hub, port: port.poe_mode != "off",
name_fn=lambda port: f"{port.name} PoE", name_fn=lambda port: f"{port.name} PoE",
object_fn=lambda api, obj_id: api.ports[obj_id], object_fn=lambda api, obj_id: api.ports[obj_id],
should_poll=False, should_poll=False,
supported_fn=lambda controller, obj_id: controller.api.ports[obj_id].port_poe, supported_fn=lambda hub, obj_id: hub.api.ports[obj_id].port_poe,
unique_id_fn=lambda controller, obj_id: f"poe-{obj_id}", unique_id_fn=lambda hub, obj_id: f"poe-{obj_id}",
), ),
UnifiSwitchEntityDescription[Wlans, Wlan]( UnifiSwitchEntityDescription[Wlans, Wlan](
key="WLAN control", key="WLAN control",
@ -306,19 +288,19 @@ ENTITY_DESCRIPTIONS: tuple[UnifiSwitchEntityDescription, ...] = (
entity_category=EntityCategory.CONFIG, entity_category=EntityCategory.CONFIG,
has_entity_name=True, has_entity_name=True,
icon="mdi:wifi-check", icon="mdi:wifi-check",
allowed_fn=lambda controller, obj_id: True, allowed_fn=lambda hub, obj_id: True,
api_handler_fn=lambda api: api.wlans, api_handler_fn=lambda api: api.wlans,
available_fn=lambda controller, _: controller.available, available_fn=lambda hub, _: hub.available,
control_fn=async_wlan_control_fn, control_fn=async_wlan_control_fn,
device_info_fn=async_wlan_device_info_fn, device_info_fn=async_wlan_device_info_fn,
event_is_on=None, event_is_on=None,
event_to_subscribe=None, event_to_subscribe=None,
is_on_fn=lambda controller, wlan: wlan.enabled, is_on_fn=lambda hub, wlan: wlan.enabled,
name_fn=lambda wlan: None, name_fn=lambda wlan: None,
object_fn=lambda api, obj_id: api.wlans[obj_id], object_fn=lambda api, obj_id: api.wlans[obj_id],
should_poll=False, should_poll=False,
supported_fn=lambda controller, obj_id: True, supported_fn=lambda hub, obj_id: True,
unique_id_fn=lambda controller, obj_id: f"wlan-{obj_id}", unique_id_fn=lambda hub, obj_id: f"wlan-{obj_id}",
), ),
) )
@ -329,7 +311,7 @@ def async_update_unique_id(hass: HomeAssistant, config_entry: ConfigEntry) -> No
Introduced with release 2023.12. Introduced with release 2023.12.
""" """
controller: UniFiController = hass.data[UNIFI_DOMAIN][config_entry.entry_id] hub: UnifiHub = hass.data[UNIFI_DOMAIN][config_entry.entry_id]
ent_reg = er.async_get(hass) ent_reg = er.async_get(hass)
@callback @callback
@ -344,10 +326,10 @@ def async_update_unique_id(hass: HomeAssistant, config_entry: ConfigEntry) -> No
if entity_id := ent_reg.async_get_entity_id(DOMAIN, UNIFI_DOMAIN, unique_id): if entity_id := ent_reg.async_get_entity_id(DOMAIN, UNIFI_DOMAIN, unique_id):
ent_reg.async_update_entity(entity_id, new_unique_id=new_unique_id) ent_reg.async_update_entity(entity_id, new_unique_id=new_unique_id)
for obj_id in controller.api.outlets: for obj_id in hub.api.outlets:
update_unique_id(obj_id, "outlet") update_unique_id(obj_id, "outlet")
for obj_id in controller.api.ports: for obj_id in hub.api.ports:
update_unique_id(obj_id, "poe") update_unique_id(obj_id, "poe")
@ -358,7 +340,7 @@ async def async_setup_entry(
) -> None: ) -> None:
"""Set up switches for UniFi Network integration.""" """Set up switches for UniFi Network integration."""
async_update_unique_id(hass, config_entry) async_update_unique_id(hass, config_entry)
UniFiController.register_platform( UnifiHub.register_platform(
hass, hass,
config_entry, config_entry,
async_add_entities, async_add_entities,
@ -384,11 +366,11 @@ class UnifiSwitchEntity(UnifiEntity[HandlerT, ApiItemT], SwitchEntity):
async def async_turn_on(self, **kwargs: Any) -> None: async def async_turn_on(self, **kwargs: Any) -> None:
"""Turn on switch.""" """Turn on switch."""
await self.entity_description.control_fn(self.controller, self._obj_id, True) await self.entity_description.control_fn(self.hub, self._obj_id, True)
async def async_turn_off(self, **kwargs: Any) -> None: async def async_turn_off(self, **kwargs: Any) -> None:
"""Turn off switch.""" """Turn off switch."""
await self.entity_description.control_fn(self.controller, self._obj_id, False) await self.entity_description.control_fn(self.hub, self._obj_id, False)
@callback @callback
def async_update_state(self, event: ItemEvent, obj_id: str) -> None: def async_update_state(self, event: ItemEvent, obj_id: str) -> None:
@ -400,8 +382,8 @@ class UnifiSwitchEntity(UnifiEntity[HandlerT, ApiItemT], SwitchEntity):
return return
description = self.entity_description description = self.entity_description
obj = description.object_fn(self.controller.api, self._obj_id) obj = description.object_fn(self.hub.api, self._obj_id)
if (is_on := description.is_on_fn(self.controller, obj)) != self.is_on: if (is_on := description.is_on_fn(self.hub, obj)) != self.is_on:
self._attr_is_on = is_on self._attr_is_on = is_on
@callback @callback
@ -416,7 +398,7 @@ class UnifiSwitchEntity(UnifiEntity[HandlerT, ApiItemT], SwitchEntity):
if event.key in description.event_to_subscribe: if event.key in description.event_to_subscribe:
self._attr_is_on = event.key in description.event_is_on self._attr_is_on = event.key in description.event_is_on
self._attr_available = description.available_fn(self.controller, self._obj_id) self._attr_available = description.available_fn(self.hub, self._obj_id)
self.async_write_ha_state() self.async_write_ha_state()
async def async_added_to_hass(self) -> None: async def async_added_to_hass(self) -> None:
@ -425,7 +407,7 @@ class UnifiSwitchEntity(UnifiEntity[HandlerT, ApiItemT], SwitchEntity):
if self.entity_description.custom_subscribe is not None: if self.entity_description.custom_subscribe is not None:
self.async_on_remove( self.async_on_remove(
self.entity_description.custom_subscribe(self.controller.api)( self.entity_description.custom_subscribe(self.hub.api)(
self.async_signalling_callback, ItemEvent.CHANGED self.async_signalling_callback, ItemEvent.CHANGED
), ),
) )

View file

@ -21,13 +21,13 @@ from homeassistant.config_entries import ConfigEntry
from homeassistant.core import HomeAssistant, callback from homeassistant.core import HomeAssistant, callback
from homeassistant.helpers.entity_platform import AddEntitiesCallback from homeassistant.helpers.entity_platform import AddEntitiesCallback
from .controller import UniFiController
from .entity import ( from .entity import (
UnifiEntity, UnifiEntity,
UnifiEntityDescription, UnifiEntityDescription,
async_device_available_fn, async_device_available_fn,
async_device_device_info_fn, async_device_device_info_fn,
) )
from .hub import UnifiHub
LOGGER = logging.getLogger(__name__) LOGGER = logging.getLogger(__name__)
@ -62,7 +62,7 @@ ENTITY_DESCRIPTIONS: tuple[UnifiUpdateEntityDescription, ...] = (
key="Upgrade device", key="Upgrade device",
device_class=UpdateDeviceClass.FIRMWARE, device_class=UpdateDeviceClass.FIRMWARE,
has_entity_name=True, has_entity_name=True,
allowed_fn=lambda controller, obj_id: True, allowed_fn=lambda hub, obj_id: True,
api_handler_fn=lambda api: api.devices, api_handler_fn=lambda api: api.devices,
available_fn=async_device_available_fn, available_fn=async_device_available_fn,
control_fn=async_device_control_fn, control_fn=async_device_control_fn,
@ -73,8 +73,8 @@ ENTITY_DESCRIPTIONS: tuple[UnifiUpdateEntityDescription, ...] = (
object_fn=lambda api, obj_id: api.devices[obj_id], object_fn=lambda api, obj_id: api.devices[obj_id],
should_poll=False, should_poll=False,
state_fn=lambda api, device: device.state == 4, state_fn=lambda api, device: device.state == 4,
supported_fn=lambda controller, obj_id: True, supported_fn=lambda hub, obj_id: True,
unique_id_fn=lambda controller, obj_id: f"device_update-{obj_id}", unique_id_fn=lambda hub, obj_id: f"device_update-{obj_id}",
), ),
) )
@ -85,7 +85,7 @@ async def async_setup_entry(
async_add_entities: AddEntitiesCallback, async_add_entities: AddEntitiesCallback,
) -> None: ) -> None:
"""Set up update entities for UniFi Network integration.""" """Set up update entities for UniFi Network integration."""
UniFiController.register_platform( UnifiHub.register_platform(
hass, hass,
config_entry, config_entry,
async_add_entities, async_add_entities,
@ -103,7 +103,7 @@ class UnifiDeviceUpdateEntity(UnifiEntity[_HandlerT, _DataT], UpdateEntity):
def async_initiate_state(self) -> None: def async_initiate_state(self) -> None:
"""Initiate entity state.""" """Initiate entity state."""
self._attr_supported_features = UpdateEntityFeature.PROGRESS self._attr_supported_features = UpdateEntityFeature.PROGRESS
if self.controller.is_admin: if self.hub.is_admin:
self._attr_supported_features |= UpdateEntityFeature.INSTALL self._attr_supported_features |= UpdateEntityFeature.INSTALL
self.async_update_state(ItemEvent.ADDED, self._obj_id) self.async_update_state(ItemEvent.ADDED, self._obj_id)
@ -112,7 +112,7 @@ class UnifiDeviceUpdateEntity(UnifiEntity[_HandlerT, _DataT], UpdateEntity):
self, version: str | None, backup: bool, **kwargs: Any self, version: str | None, backup: bool, **kwargs: Any
) -> None: ) -> None:
"""Install an update.""" """Install an update."""
await self.entity_description.control_fn(self.controller.api, self._obj_id) await self.entity_description.control_fn(self.hub.api, self._obj_id)
@callback @callback
def async_update_state(self, event: ItemEvent, obj_id: str) -> None: def async_update_state(self, event: ItemEvent, obj_id: str) -> None:
@ -122,7 +122,7 @@ class UnifiDeviceUpdateEntity(UnifiEntity[_HandlerT, _DataT], UpdateEntity):
""" """
description = self.entity_description description = self.entity_description
obj = description.object_fn(self.controller.api, self._obj_id) obj = description.object_fn(self.hub.api, self._obj_id)
self._attr_in_progress = description.state_fn(self.controller.api, obj) self._attr_in_progress = description.state_fn(self.hub.api, obj)
self._attr_installed_version = obj.version self._attr_installed_version = obj.version
self._attr_latest_version = obj.upgrade_to_firmware or obj.version self._attr_latest_version = obj.upgrade_to_firmware or obj.version

View file

@ -9,14 +9,14 @@ from aiounifi.models.message import MessageKey
import pytest import pytest
from homeassistant.components.unifi.const import DOMAIN as UNIFI_DOMAIN from homeassistant.components.unifi.const import DOMAIN as UNIFI_DOMAIN
from homeassistant.components.unifi.controller import RETRY_TIMER from homeassistant.components.unifi.hub import RETRY_TIMER
from homeassistant.const import CONTENT_TYPE_JSON from homeassistant.const import CONTENT_TYPE_JSON
from homeassistant.core import HomeAssistant from homeassistant.core import HomeAssistant
from homeassistant.helpers import device_registry as dr from homeassistant.helpers import device_registry as dr
import homeassistant.util.dt as dt_util import homeassistant.util.dt as dt_util
from tests.common import MockConfigEntry, async_fire_time_changed from tests.common import MockConfigEntry, async_fire_time_changed
from tests.components.unifi.test_controller import DEFAULT_CONFIG_ENTRY_ID from tests.components.unifi.test_hub import DEFAULT_CONFIG_ENTRY_ID
from tests.test_util.aiohttp import AiohttpClientMocker from tests.test_util.aiohttp import AiohttpClientMocker
@ -43,12 +43,12 @@ class WebsocketStateManager(asyncio.Event):
Mock api calls done by 'await self.api.login'. Mock api calls done by 'await self.api.login'.
Fail will make 'await self.api.start_websocket' return immediately. Fail will make 'await self.api.start_websocket' return immediately.
""" """
controller = self.hass.data[UNIFI_DOMAIN][DEFAULT_CONFIG_ENTRY_ID] hub = self.hass.data[UNIFI_DOMAIN][DEFAULT_CONFIG_ENTRY_ID]
self.aioclient_mock.get( self.aioclient_mock.get(
f"https://{controller.host}:1234", status=302 f"https://{hub.host}:1234", status=302
) # Check UniFi OS ) # Check UniFi OS
self.aioclient_mock.post( self.aioclient_mock.post(
f"https://{controller.host}:1234/api/login", f"https://{hub.host}:1234/api/login",
json={"data": "login successful", "meta": {"rc": "ok"}}, json={"data": "login successful", "meta": {"rc": "ok"}},
headers={"content-type": CONTENT_TYPE_JSON}, headers={"content-type": CONTENT_TYPE_JSON},
) )
@ -79,13 +79,13 @@ def mock_unifi_websocket(hass):
data: list[dict] | dict | None = None, data: list[dict] | dict | None = None,
): ):
"""Generate a websocket call.""" """Generate a websocket call."""
controller = hass.data[UNIFI_DOMAIN][DEFAULT_CONFIG_ENTRY_ID] hub = hass.data[UNIFI_DOMAIN][DEFAULT_CONFIG_ENTRY_ID]
if data and not message: if data and not message:
controller.api.messages.handler(data) hub.api.messages.handler(data)
elif data and message: elif data and message:
if not isinstance(data, list): if not isinstance(data, list):
data = [data] data = [data]
controller.api.messages.handler( hub.api.messages.handler(
{ {
"meta": {"message": message.value}, "meta": {"message": message.value},
"data": data, "data": data,

View file

@ -6,7 +6,7 @@ from homeassistant.const import ATTR_DEVICE_CLASS, STATE_UNAVAILABLE, EntityCate
from homeassistant.core import HomeAssistant from homeassistant.core import HomeAssistant
from homeassistant.helpers import entity_registry as er from homeassistant.helpers import entity_registry as er
from .test_controller import setup_unifi_integration from .test_hub import setup_unifi_integration
from tests.test_util.aiohttp import AiohttpClientMocker from tests.test_util.aiohttp import AiohttpClientMocker
@ -36,7 +36,7 @@ async def test_restart_device_button(
} }
], ],
) )
controller = hass.data[UNIFI_DOMAIN][config_entry.entry_id] hub = hass.data[UNIFI_DOMAIN][config_entry.entry_id]
assert len(hass.states.async_entity_ids(BUTTON_DOMAIN)) == 1 assert len(hass.states.async_entity_ids(BUTTON_DOMAIN)) == 1
@ -52,7 +52,7 @@ async def test_restart_device_button(
# Send restart device command # Send restart device command
aioclient_mock.clear_requests() aioclient_mock.clear_requests()
aioclient_mock.post( aioclient_mock.post(
f"https://{controller.host}:1234/api/s/{controller.site}/cmd/devmgr", f"https://{hub.host}:1234/api/s/{hub.site}/cmd/devmgr",
) )
await hass.services.async_call( await hass.services.async_call(
@ -120,7 +120,7 @@ async def test_power_cycle_poe(
} }
], ],
) )
controller = hass.data[UNIFI_DOMAIN][config_entry.entry_id] hub = hass.data[UNIFI_DOMAIN][config_entry.entry_id]
assert len(hass.states.async_entity_ids(BUTTON_DOMAIN)) == 2 assert len(hass.states.async_entity_ids(BUTTON_DOMAIN)) == 2
@ -136,7 +136,7 @@ async def test_power_cycle_poe(
# Send restart device command # Send restart device command
aioclient_mock.clear_requests() aioclient_mock.clear_requests()
aioclient_mock.post( aioclient_mock.post(
f"https://{controller.host}:1234/api/s/{controller.site}/cmd/devmgr", f"https://{hub.host}:1234/api/s/{hub.site}/cmd/devmgr",
) )
await hass.services.async_call( await hass.services.async_call(

View file

@ -33,7 +33,7 @@ from homeassistant.const import (
) )
from homeassistant.core import HomeAssistant from homeassistant.core import HomeAssistant
from .test_controller import setup_unifi_integration from .test_hub import setup_unifi_integration
from tests.common import MockConfigEntry from tests.common import MockConfigEntry
from tests.test_util.aiohttp import AiohttpClientMocker from tests.test_util.aiohttp import AiohttpClientMocker
@ -356,7 +356,7 @@ async def test_flow_fails_user_credentials_faulty(
assert result["errors"] == {"base": "faulty_credentials"} assert result["errors"] == {"base": "faulty_credentials"}
async def test_flow_fails_controller_unavailable( async def test_flow_fails_hub_unavailable(
hass: HomeAssistant, aioclient_mock: AiohttpClientMocker hass: HomeAssistant, aioclient_mock: AiohttpClientMocker
) -> None: ) -> None:
"""Test config flow.""" """Test config flow."""
@ -388,10 +388,10 @@ async def test_flow_fails_controller_unavailable(
async def test_reauth_flow_update_configuration( async def test_reauth_flow_update_configuration(
hass: HomeAssistant, aioclient_mock: AiohttpClientMocker hass: HomeAssistant, aioclient_mock: AiohttpClientMocker
) -> None: ) -> None:
"""Verify reauth flow can update controller configuration.""" """Verify reauth flow can update hub configuration."""
config_entry = await setup_unifi_integration(hass, aioclient_mock) config_entry = await setup_unifi_integration(hass, aioclient_mock)
controller = hass.data[UNIFI_DOMAIN][config_entry.entry_id] hub = hass.data[UNIFI_DOMAIN][config_entry.entry_id]
controller.available = False hub.available = False
result = await hass.config_entries.flow.async_init( result = await hass.config_entries.flow.async_init(
UNIFI_DOMAIN, UNIFI_DOMAIN,

View file

@ -21,7 +21,7 @@ from homeassistant.core import HomeAssistant
from homeassistant.helpers import entity_registry as er from homeassistant.helpers import entity_registry as er
import homeassistant.util.dt as dt_util import homeassistant.util.dt as dt_util
from .test_controller import ENTRY_CONFIG, setup_unifi_integration from .test_hub import ENTRY_CONFIG, setup_unifi_integration
from tests.common import async_fire_time_changed from tests.common import async_fire_time_changed
from tests.test_util.aiohttp import AiohttpClientMocker from tests.test_util.aiohttp import AiohttpClientMocker
@ -55,7 +55,7 @@ async def test_tracked_wireless_clients(
config_entry = await setup_unifi_integration( config_entry = await setup_unifi_integration(
hass, aioclient_mock, clients_response=[client] hass, aioclient_mock, clients_response=[client]
) )
controller = hass.data[UNIFI_DOMAIN][config_entry.entry_id] hub = hass.data[UNIFI_DOMAIN][config_entry.entry_id]
assert len(hass.states.async_entity_ids(TRACKER_DOMAIN)) == 1 assert len(hass.states.async_entity_ids(TRACKER_DOMAIN)) == 1
assert hass.states.get("device_tracker.client").state == STATE_NOT_HOME assert hass.states.get("device_tracker.client").state == STATE_NOT_HOME
@ -70,7 +70,7 @@ async def test_tracked_wireless_clients(
# Change time to mark client as away # Change time to mark client as away
new_time = dt_util.utcnow() + controller.option_detection_time new_time = dt_util.utcnow() + hub.option_detection_time
with freeze_time(new_time): with freeze_time(new_time):
async_fire_time_changed(hass, new_time) async_fire_time_changed(hass, new_time)
await hass.async_block_till_done() await hass.async_block_till_done()
@ -194,7 +194,7 @@ async def test_tracked_wireless_clients_event_source(
config_entry = await setup_unifi_integration( config_entry = await setup_unifi_integration(
hass, aioclient_mock, clients_response=[client] hass, aioclient_mock, clients_response=[client]
) )
controller = hass.data[UNIFI_DOMAIN][config_entry.entry_id] hub = hass.data[UNIFI_DOMAIN][config_entry.entry_id]
assert len(hass.states.async_entity_ids(TRACKER_DOMAIN)) == 1 assert len(hass.states.async_entity_ids(TRACKER_DOMAIN)) == 1
assert hass.states.get("device_tracker.client").state == STATE_NOT_HOME assert hass.states.get("device_tracker.client").state == STATE_NOT_HOME
@ -243,7 +243,7 @@ async def test_tracked_wireless_clients_event_source(
assert hass.states.get("device_tracker.client").state == STATE_HOME assert hass.states.get("device_tracker.client").state == STATE_HOME
# Change time to mark client as away # Change time to mark client as away
freezer.tick(controller.option_detection_time + timedelta(seconds=1)) freezer.tick(hub.option_detection_time + timedelta(seconds=1))
async_fire_time_changed(hass) async_fire_time_changed(hass)
await hass.async_block_till_done() await hass.async_block_till_done()
@ -282,7 +282,7 @@ async def test_tracked_wireless_clients_event_source(
assert hass.states.get("device_tracker.client").state == STATE_HOME assert hass.states.get("device_tracker.client").state == STATE_HOME
# Change time to mark client as away # Change time to mark client as away
freezer.tick(controller.option_detection_time + timedelta(seconds=1)) freezer.tick(hub.option_detection_time + timedelta(seconds=1))
async_fire_time_changed(hass) async_fire_time_changed(hass)
await hass.async_block_till_done() await hass.async_block_till_done()
@ -407,13 +407,13 @@ async def test_remove_clients(
assert hass.states.get("device_tracker.client_2") assert hass.states.get("device_tracker.client_2")
async def test_controller_state_change( async def test_hub_state_change(
hass: HomeAssistant, hass: HomeAssistant,
aioclient_mock: AiohttpClientMocker, aioclient_mock: AiohttpClientMocker,
websocket_mock, websocket_mock,
mock_device_registry, mock_device_registry,
) -> None: ) -> None:
"""Verify entities state reflect on controller becoming unavailable.""" """Verify entities state reflect on hub connection becoming unavailable."""
client = { client = {
"essid": "ssid", "essid": "ssid",
"hostname": "client", "hostname": "client",
@ -682,7 +682,7 @@ async def test_option_ssid_filter(
config_entry = await setup_unifi_integration( config_entry = await setup_unifi_integration(
hass, aioclient_mock, clients_response=[client, client_on_ssid2] hass, aioclient_mock, clients_response=[client, client_on_ssid2]
) )
controller = hass.data[UNIFI_DOMAIN][config_entry.entry_id] hub = hass.data[UNIFI_DOMAIN][config_entry.entry_id]
assert len(hass.states.async_entity_ids(TRACKER_DOMAIN)) == 2 assert len(hass.states.async_entity_ids(TRACKER_DOMAIN)) == 2
@ -711,7 +711,7 @@ async def test_option_ssid_filter(
mock_unifi_websocket(message=MessageKey.CLIENT, data=client_on_ssid2) mock_unifi_websocket(message=MessageKey.CLIENT, data=client_on_ssid2)
await hass.async_block_till_done() await hass.async_block_till_done()
new_time = dt_util.utcnow() + controller.option_detection_time new_time = dt_util.utcnow() + hub.option_detection_time
with freeze_time(new_time): with freeze_time(new_time):
async_fire_time_changed(hass, new_time) async_fire_time_changed(hass, new_time)
await hass.async_block_till_done() await hass.async_block_till_done()
@ -739,7 +739,7 @@ async def test_option_ssid_filter(
# Time pass to mark client as away # Time pass to mark client as away
new_time += controller.option_detection_time new_time += hub.option_detection_time
with freeze_time(new_time): with freeze_time(new_time):
async_fire_time_changed(hass, new_time) async_fire_time_changed(hass, new_time)
await hass.async_block_till_done() await hass.async_block_till_done()
@ -758,7 +758,7 @@ async def test_option_ssid_filter(
mock_unifi_websocket(message=MessageKey.CLIENT, data=client_on_ssid2) mock_unifi_websocket(message=MessageKey.CLIENT, data=client_on_ssid2)
await hass.async_block_till_done() await hass.async_block_till_done()
new_time += controller.option_detection_time new_time += hub.option_detection_time
with freeze_time(new_time): with freeze_time(new_time):
async_fire_time_changed(hass, new_time) async_fire_time_changed(hass, new_time)
await hass.async_block_till_done() await hass.async_block_till_done()
@ -788,7 +788,7 @@ async def test_wireless_client_go_wired_issue(
config_entry = await setup_unifi_integration( config_entry = await setup_unifi_integration(
hass, aioclient_mock, clients_response=[client] hass, aioclient_mock, clients_response=[client]
) )
controller = hass.data[UNIFI_DOMAIN][config_entry.entry_id] hub = hass.data[UNIFI_DOMAIN][config_entry.entry_id]
assert len(hass.states.async_entity_ids(TRACKER_DOMAIN)) == 1 assert len(hass.states.async_entity_ids(TRACKER_DOMAIN)) == 1
@ -807,7 +807,7 @@ async def test_wireless_client_go_wired_issue(
assert client_state.state == STATE_HOME assert client_state.state == STATE_HOME
# Pass time # Pass time
new_time = dt_util.utcnow() + controller.option_detection_time new_time = dt_util.utcnow() + hub.option_detection_time
with freeze_time(new_time): with freeze_time(new_time):
async_fire_time_changed(hass, new_time) async_fire_time_changed(hass, new_time)
await hass.async_block_till_done() await hass.async_block_till_done()
@ -859,7 +859,7 @@ async def test_option_ignore_wired_bug(
options={CONF_IGNORE_WIRED_BUG: True}, options={CONF_IGNORE_WIRED_BUG: True},
clients_response=[client], clients_response=[client],
) )
controller = hass.data[UNIFI_DOMAIN][config_entry.entry_id] hub = hass.data[UNIFI_DOMAIN][config_entry.entry_id]
assert len(hass.states.async_entity_ids(TRACKER_DOMAIN)) == 1 assert len(hass.states.async_entity_ids(TRACKER_DOMAIN)) == 1
# Client is wireless # Client is wireless
@ -876,7 +876,7 @@ async def test_option_ignore_wired_bug(
assert client_state.state == STATE_HOME assert client_state.state == STATE_HOME
# pass time # pass time
new_time = dt_util.utcnow() + controller.option_detection_time new_time = dt_util.utcnow() + hub.option_detection_time
with freeze_time(new_time): with freeze_time(new_time):
async_fire_time_changed(hass, new_time) async_fire_time_changed(hass, new_time)
await hass.async_block_till_done() await hass.async_block_till_done()

View file

@ -7,7 +7,7 @@ from homeassistant.components.unifi.const import (
) )
from homeassistant.core import HomeAssistant from homeassistant.core import HomeAssistant
from .test_controller import setup_unifi_integration from .test_hub import setup_unifi_integration
from tests.components.diagnostics import get_diagnostics_for_config_entry from tests.components.diagnostics import get_diagnostics_for_config_entry
from tests.test_util.aiohttp import AiohttpClientMocker from tests.test_util.aiohttp import AiohttpClientMocker

View file

@ -26,8 +26,8 @@ from homeassistant.components.unifi.const import (
PLATFORMS, PLATFORMS,
UNIFI_WIRELESS_CLIENTS, UNIFI_WIRELESS_CLIENTS,
) )
from homeassistant.components.unifi.controller import get_unifi_controller
from homeassistant.components.unifi.errors import AuthenticationRequired, CannotConnect from homeassistant.components.unifi.errors import AuthenticationRequired, CannotConnect
from homeassistant.components.unifi.hub import get_unifi_api
from homeassistant.const import ( from homeassistant.const import (
CONF_HOST, CONF_HOST,
CONF_PASSWORD, CONF_PASSWORD,
@ -194,7 +194,6 @@ async def setup_unifi_integration(
system_information_response=None, system_information_response=None,
wlans_response=None, wlans_response=None,
known_wireless_clients=None, known_wireless_clients=None,
controllers=None,
unique_id="1", unique_id="1",
config_entry_id=DEFAULT_CONFIG_ENTRY_ID, config_entry_id=DEFAULT_CONFIG_ENTRY_ID,
): ):
@ -241,7 +240,7 @@ async def setup_unifi_integration(
return config_entry return config_entry
async def test_controller_setup( async def test_hub_setup(
hass: HomeAssistant, hass: HomeAssistant,
device_registry: dr.DeviceRegistry, device_registry: dr.DeviceRegistry,
aioclient_mock: AiohttpClientMocker, aioclient_mock: AiohttpClientMocker,
@ -254,9 +253,9 @@ async def test_controller_setup(
config_entry = await setup_unifi_integration( config_entry = await setup_unifi_integration(
hass, aioclient_mock, system_information_response=SYSTEM_INFORMATION hass, aioclient_mock, system_information_response=SYSTEM_INFORMATION
) )
controller = hass.data[UNIFI_DOMAIN][config_entry.entry_id] hub = hass.data[UNIFI_DOMAIN][config_entry.entry_id]
entry = controller.config_entry entry = hub.config_entry
assert len(forward_entry_setup.mock_calls) == len(PLATFORMS) 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[0][1] == (entry, BUTTON_DOMAIN)
assert forward_entry_setup.mock_calls[1][1] == (entry, TRACKER_DOMAIN) assert forward_entry_setup.mock_calls[1][1] == (entry, TRACKER_DOMAIN)
@ -264,21 +263,21 @@ async def test_controller_setup(
assert forward_entry_setup.mock_calls[3][1] == (entry, SENSOR_DOMAIN) assert forward_entry_setup.mock_calls[3][1] == (entry, SENSOR_DOMAIN)
assert forward_entry_setup.mock_calls[4][1] == (entry, SWITCH_DOMAIN) assert forward_entry_setup.mock_calls[4][1] == (entry, SWITCH_DOMAIN)
assert controller.host == ENTRY_CONFIG[CONF_HOST] assert hub.host == ENTRY_CONFIG[CONF_HOST]
assert controller.is_admin == (SITE[0]["role"] == "admin") assert hub.is_admin == (SITE[0]["role"] == "admin")
assert controller.option_allow_bandwidth_sensors == DEFAULT_ALLOW_BANDWIDTH_SENSORS assert hub.option_allow_bandwidth_sensors == DEFAULT_ALLOW_BANDWIDTH_SENSORS
assert controller.option_allow_uptime_sensors == DEFAULT_ALLOW_UPTIME_SENSORS assert hub.option_allow_uptime_sensors == DEFAULT_ALLOW_UPTIME_SENSORS
assert isinstance(controller.option_block_clients, list) assert isinstance(hub.option_block_clients, list)
assert controller.option_track_clients == DEFAULT_TRACK_CLIENTS assert hub.option_track_clients == DEFAULT_TRACK_CLIENTS
assert controller.option_track_devices == DEFAULT_TRACK_DEVICES assert hub.option_track_devices == DEFAULT_TRACK_DEVICES
assert controller.option_track_wired_clients == DEFAULT_TRACK_WIRED_CLIENTS assert hub.option_track_wired_clients == DEFAULT_TRACK_WIRED_CLIENTS
assert controller.option_detection_time == timedelta(seconds=DEFAULT_DETECTION_TIME) assert hub.option_detection_time == timedelta(seconds=DEFAULT_DETECTION_TIME)
assert isinstance(controller.option_ssid_filter, set) assert isinstance(hub.option_ssid_filter, set)
assert controller.signal_reachable == "unifi-reachable-1" assert hub.signal_reachable == "unifi-reachable-1"
assert controller.signal_options_update == "unifi-options-1" assert hub.signal_options_update == "unifi-options-1"
assert controller.signal_heartbeat_missed == "unifi-heartbeat-missed" assert hub.signal_heartbeat_missed == "unifi-heartbeat-missed"
device_entry = device_registry.async_get_or_create( device_entry = device_registry.async_get_or_create(
config_entry_id=config_entry.entry_id, config_entry_id=config_entry.entry_id,
@ -288,20 +287,20 @@ async def test_controller_setup(
assert device_entry.sw_version == "7.4.162" assert device_entry.sw_version == "7.4.162"
async def test_controller_not_accessible(hass: HomeAssistant) -> None: async def test_hub_not_accessible(hass: HomeAssistant) -> None:
"""Retry to login gets scheduled when connection fails.""" """Retry to login gets scheduled when connection fails."""
with patch( with patch(
"homeassistant.components.unifi.controller.get_unifi_controller", "homeassistant.components.unifi.hub.get_unifi_api",
side_effect=CannotConnect, side_effect=CannotConnect,
): ):
await setup_unifi_integration(hass) await setup_unifi_integration(hass)
assert hass.data[UNIFI_DOMAIN] == {} assert hass.data[UNIFI_DOMAIN] == {}
async def test_controller_trigger_reauth_flow(hass: HomeAssistant) -> None: async def test_hub_trigger_reauth_flow(hass: HomeAssistant) -> None:
"""Failed authentication trigger a reauthentication flow.""" """Failed authentication trigger a reauthentication flow."""
with patch( with patch(
"homeassistant.components.unifi.get_unifi_controller", "homeassistant.components.unifi.get_unifi_api",
side_effect=AuthenticationRequired, side_effect=AuthenticationRequired,
), patch.object(hass.config_entries.flow, "async_init") as mock_flow_init: ), patch.object(hass.config_entries.flow, "async_init") as mock_flow_init:
await setup_unifi_integration(hass) await setup_unifi_integration(hass)
@ -309,10 +308,10 @@ async def test_controller_trigger_reauth_flow(hass: HomeAssistant) -> None:
assert hass.data[UNIFI_DOMAIN] == {} assert hass.data[UNIFI_DOMAIN] == {}
async def test_controller_unknown_error(hass: HomeAssistant) -> None: async def test_hub_unknown_error(hass: HomeAssistant) -> None:
"""Unknown errors are handled.""" """Unknown errors are handled."""
with patch( with patch(
"homeassistant.components.unifi.controller.get_unifi_controller", "homeassistant.components.unifi.hub.get_unifi_api",
side_effect=Exception, side_effect=Exception,
): ):
await setup_unifi_integration(hass) await setup_unifi_integration(hass)
@ -324,10 +323,10 @@ async def test_config_entry_updated(
) -> None: ) -> None:
"""Calling reset when the entry has been setup.""" """Calling reset when the entry has been setup."""
config_entry = await setup_unifi_integration(hass, aioclient_mock) config_entry = await setup_unifi_integration(hass, aioclient_mock)
controller = hass.data[UNIFI_DOMAIN][config_entry.entry_id] hub = hass.data[UNIFI_DOMAIN][config_entry.entry_id]
event_call = Mock() event_call = Mock()
unsub = async_dispatcher_connect(hass, controller.signal_options_update, event_call) unsub = async_dispatcher_connect(hass, hub.signal_options_update, event_call)
hass.config_entries.async_update_entry( hass.config_entries.async_update_entry(
config_entry, options={CONF_TRACK_CLIENTS: False, CONF_TRACK_DEVICES: False} config_entry, options={CONF_TRACK_CLIENTS: False, CONF_TRACK_DEVICES: False}
@ -347,9 +346,9 @@ async def test_reset_after_successful_setup(
) -> None: ) -> None:
"""Calling reset when the entry has been setup.""" """Calling reset when the entry has been setup."""
config_entry = await setup_unifi_integration(hass, aioclient_mock) config_entry = await setup_unifi_integration(hass, aioclient_mock)
controller = hass.data[UNIFI_DOMAIN][config_entry.entry_id] hub = hass.data[UNIFI_DOMAIN][config_entry.entry_id]
result = await controller.async_reset() result = await hub.async_reset()
await hass.async_block_till_done() await hass.async_block_till_done()
assert result is True assert result is True
@ -360,13 +359,13 @@ async def test_reset_fails(
) -> None: ) -> None:
"""Calling reset when the entry has been setup can return false.""" """Calling reset when the entry has been setup can return false."""
config_entry = await setup_unifi_integration(hass, aioclient_mock) config_entry = await setup_unifi_integration(hass, aioclient_mock)
controller = hass.data[UNIFI_DOMAIN][config_entry.entry_id] hub = hass.data[UNIFI_DOMAIN][config_entry.entry_id]
with patch( with patch(
"homeassistant.config_entries.ConfigEntries.async_forward_entry_unload", "homeassistant.config_entries.ConfigEntries.async_forward_entry_unload",
return_value=False, return_value=False,
): ):
result = await controller.async_reset() result = await hub.async_reset()
await hass.async_block_till_done() await hass.async_block_till_done()
assert result is False assert result is False
@ -435,7 +434,7 @@ async def test_reconnect_mechanism_exceptions(
await setup_unifi_integration(hass, aioclient_mock) await setup_unifi_integration(hass, aioclient_mock)
with patch("aiounifi.Controller.login", side_effect=exception), patch( with patch("aiounifi.Controller.login", side_effect=exception), patch(
"homeassistant.components.unifi.controller.UniFiController.reconnect" "homeassistant.components.unifi.hub.UnifiHub.reconnect"
) as mock_reconnect: ) as mock_reconnect:
await websocket_mock.disconnect() await websocket_mock.disconnect()
@ -443,18 +442,18 @@ async def test_reconnect_mechanism_exceptions(
mock_reconnect.assert_called_once() mock_reconnect.assert_called_once()
async def test_get_unifi_controller(hass: HomeAssistant) -> None: async def test_get_unifi_api(hass: HomeAssistant) -> None:
"""Successful call.""" """Successful call."""
with patch("aiounifi.Controller.login", return_value=True): with patch("aiounifi.Controller.login", return_value=True):
assert await get_unifi_controller(hass, ENTRY_CONFIG) assert await get_unifi_api(hass, ENTRY_CONFIG)
async def test_get_unifi_controller_verify_ssl_false(hass: HomeAssistant) -> None: async def test_get_unifi_api_verify_ssl_false(hass: HomeAssistant) -> None:
"""Successful call with verify ssl set to false.""" """Successful call with verify ssl set to false."""
controller_data = dict(ENTRY_CONFIG) hub_data = dict(ENTRY_CONFIG)
controller_data[CONF_VERIFY_SSL] = False hub_data[CONF_VERIFY_SSL] = False
with patch("aiounifi.Controller.login", return_value=True): with patch("aiounifi.Controller.login", return_value=True):
assert await get_unifi_controller(hass, controller_data) assert await get_unifi_api(hass, hub_data)
@pytest.mark.parametrize( @pytest.mark.parametrize(
@ -471,11 +470,11 @@ async def test_get_unifi_controller_verify_ssl_false(hass: HomeAssistant) -> Non
(aiounifi.AiounifiException, AuthenticationRequired), (aiounifi.AiounifiException, AuthenticationRequired),
], ],
) )
async def test_get_unifi_controller_fails_to_connect( async def test_get_unifi_api_fails_to_connect(
hass: HomeAssistant, side_effect, raised_exception hass: HomeAssistant, side_effect, raised_exception
) -> None: ) -> None:
"""Check that get_unifi_controller can handle controller being unavailable.""" """Check that get_unifi_api can handle UniFi Network being unavailable."""
with patch("aiounifi.Controller.login", side_effect=side_effect), pytest.raises( with patch("aiounifi.Controller.login", side_effect=side_effect), pytest.raises(
raised_exception raised_exception
): ):
await get_unifi_controller(hass, ENTRY_CONFIG) await get_unifi_api(hass, ENTRY_CONFIG)

View file

@ -15,7 +15,7 @@ from homeassistant.helpers import entity_registry as er
from homeassistant.helpers.entity_registry import RegistryEntryDisabler from homeassistant.helpers.entity_registry import RegistryEntryDisabler
from homeassistant.util import dt as dt_util from homeassistant.util import dt as dt_util
from .test_controller import setup_unifi_integration from .test_hub import setup_unifi_integration
from tests.common import async_fire_time_changed from tests.common import async_fire_time_changed
from tests.test_util.aiohttp import AiohttpClientMocker from tests.test_util.aiohttp import AiohttpClientMocker

View file

@ -8,14 +8,14 @@ from homeassistant.components.unifi.errors import AuthenticationRequired, Cannot
from homeassistant.core import HomeAssistant from homeassistant.core import HomeAssistant
from homeassistant.setup import async_setup_component from homeassistant.setup import async_setup_component
from .test_controller import DEFAULT_CONFIG_ENTRY_ID, setup_unifi_integration from .test_hub import DEFAULT_CONFIG_ENTRY_ID, setup_unifi_integration
from tests.common import flush_store from tests.common import flush_store
from tests.test_util.aiohttp import AiohttpClientMocker from tests.test_util.aiohttp import AiohttpClientMocker
async def test_setup_with_no_config(hass: HomeAssistant) -> None: async def test_setup_with_no_config(hass: HomeAssistant) -> None:
"""Test that we do not discover anything or try to set up a controller.""" """Test that we do not discover anything or try to set up a hub."""
assert await async_setup_component(hass, UNIFI_DOMAIN, {}) is True assert await async_setup_component(hass, UNIFI_DOMAIN, {}) is True
assert UNIFI_DOMAIN not in hass.data assert UNIFI_DOMAIN not in hass.data
@ -31,7 +31,7 @@ async def test_successful_config_entry(
async def test_setup_entry_fails_config_entry_not_ready(hass: HomeAssistant) -> None: async def test_setup_entry_fails_config_entry_not_ready(hass: HomeAssistant) -> None:
"""Failed authentication trigger a reauthentication flow.""" """Failed authentication trigger a reauthentication flow."""
with patch( with patch(
"homeassistant.components.unifi.get_unifi_controller", "homeassistant.components.unifi.get_unifi_api",
side_effect=CannotConnect, side_effect=CannotConnect,
): ):
await setup_unifi_integration(hass) await setup_unifi_integration(hass)
@ -42,7 +42,7 @@ async def test_setup_entry_fails_config_entry_not_ready(hass: HomeAssistant) ->
async def test_setup_entry_fails_trigger_reauth_flow(hass: HomeAssistant) -> None: async def test_setup_entry_fails_trigger_reauth_flow(hass: HomeAssistant) -> None:
"""Failed authentication trigger a reauthentication flow.""" """Failed authentication trigger a reauthentication flow."""
with patch( with patch(
"homeassistant.components.unifi.get_unifi_controller", "homeassistant.components.unifi.get_unifi_api",
side_effect=AuthenticationRequired, side_effect=AuthenticationRequired,
), patch.object(hass.config_entries.flow, "async_init") as mock_flow_init: ), patch.object(hass.config_entries.flow, "async_init") as mock_flow_init:
await setup_unifi_integration(hass) await setup_unifi_integration(hass)

View file

@ -8,7 +8,6 @@ from aiounifi.models.message import MessageKey
from freezegun.api import FrozenDateTimeFactory, freeze_time from freezegun.api import FrozenDateTimeFactory, freeze_time
import pytest import pytest
from homeassistant.components.device_tracker import DOMAIN as TRACKER_DOMAIN
from homeassistant.components.sensor import ( from homeassistant.components.sensor import (
ATTR_STATE_CLASS, ATTR_STATE_CLASS,
DOMAIN as SENSOR_DOMAIN, DOMAIN as SENSOR_DOMAIN,
@ -31,7 +30,7 @@ from homeassistant.helpers import entity_registry as er
from homeassistant.helpers.entity_registry import RegistryEntryDisabler from homeassistant.helpers.entity_registry import RegistryEntryDisabler
import homeassistant.util.dt as dt_util import homeassistant.util.dt as dt_util
from .test_controller import setup_unifi_integration from .test_hub import setup_unifi_integration
from tests.common import async_fire_time_changed from tests.common import async_fire_time_changed
from tests.test_util.aiohttp import AiohttpClientMocker from tests.test_util.aiohttp import AiohttpClientMocker
@ -396,7 +395,7 @@ async def test_bandwidth_sensors(
# Verify reset sensor after heartbeat expires # Verify reset sensor after heartbeat expires
controller = hass.data[UNIFI_DOMAIN][config_entry.entry_id] hub = hass.data[UNIFI_DOMAIN][config_entry.entry_id]
new_time = dt_util.utcnow() new_time = dt_util.utcnow()
wireless_client["last_seen"] = dt_util.as_timestamp(new_time) wireless_client["last_seen"] = dt_util.as_timestamp(new_time)
@ -410,7 +409,7 @@ async def test_bandwidth_sensors(
assert hass.states.get("sensor.wireless_client_rx").state == "3456.0" assert hass.states.get("sensor.wireless_client_rx").state == "3456.0"
assert hass.states.get("sensor.wireless_client_tx").state == "7891.0" assert hass.states.get("sensor.wireless_client_tx").state == "7891.0"
new_time = new_time + controller.option_detection_time + timedelta(seconds=1) new_time = new_time + hub.option_detection_time + timedelta(seconds=1)
with freeze_time(new_time): with freeze_time(new_time):
async_fire_time_changed(hass, new_time) async_fire_time_changed(hass, new_time)
@ -578,9 +577,7 @@ async def test_remove_sensors(
clients_response=[wired_client, wireless_client], clients_response=[wired_client, wireless_client],
) )
assert len(hass.states.async_all()) == 9
assert len(hass.states.async_entity_ids(SENSOR_DOMAIN)) == 6 assert len(hass.states.async_entity_ids(SENSOR_DOMAIN)) == 6
assert len(hass.states.async_entity_ids(TRACKER_DOMAIN)) == 2
assert hass.states.get("sensor.wired_client_rx") assert hass.states.get("sensor.wired_client_rx")
assert hass.states.get("sensor.wired_client_tx") assert hass.states.get("sensor.wired_client_tx")
assert hass.states.get("sensor.wired_client_uptime") assert hass.states.get("sensor.wired_client_uptime")
@ -593,9 +590,7 @@ async def test_remove_sensors(
mock_unifi_websocket(message=MessageKey.CLIENT_REMOVED, data=wired_client) mock_unifi_websocket(message=MessageKey.CLIENT_REMOVED, data=wired_client)
await hass.async_block_till_done() await hass.async_block_till_done()
assert len(hass.states.async_all()) == 5
assert len(hass.states.async_entity_ids(SENSOR_DOMAIN)) == 3 assert len(hass.states.async_entity_ids(SENSOR_DOMAIN)) == 3
assert len(hass.states.async_entity_ids(TRACKER_DOMAIN)) == 1
assert hass.states.get("sensor.wired_client_rx") is None assert hass.states.get("sensor.wired_client_rx") is None
assert hass.states.get("sensor.wired_client_tx") is None assert hass.states.get("sensor.wired_client_tx") is None
assert hass.states.get("sensor.wired_client_uptime") is None assert hass.states.get("sensor.wired_client_uptime") is None

View file

@ -11,7 +11,7 @@ from homeassistant.const import ATTR_DEVICE_ID
from homeassistant.core import HomeAssistant from homeassistant.core import HomeAssistant
from homeassistant.helpers import device_registry as dr from homeassistant.helpers import device_registry as dr
from .test_controller import setup_unifi_integration from .test_hub import setup_unifi_integration
from tests.test_util.aiohttp import AiohttpClientMocker from tests.test_util.aiohttp import AiohttpClientMocker
@ -66,11 +66,11 @@ async def test_reconnect_client(
config_entry = await setup_unifi_integration( config_entry = await setup_unifi_integration(
hass, aioclient_mock, clients_response=clients hass, aioclient_mock, clients_response=clients
) )
controller = hass.data[UNIFI_DOMAIN][config_entry.entry_id] hub = hass.data[UNIFI_DOMAIN][config_entry.entry_id]
aioclient_mock.clear_requests() aioclient_mock.clear_requests()
aioclient_mock.post( aioclient_mock.post(
f"https://{controller.host}:1234/api/s/{controller.site}/cmd/stamgr", f"https://{hub.host}:1234/api/s/{hub.site}/cmd/stamgr",
) )
device_entry = device_registry.async_get_or_create( device_entry = device_registry.async_get_or_create(
@ -128,12 +128,12 @@ async def test_reconnect_device_without_mac(
assert aioclient_mock.call_count == 0 assert aioclient_mock.call_count == 0
async def test_reconnect_client_controller_unavailable( async def test_reconnect_client_hub_unavailable(
hass: HomeAssistant, hass: HomeAssistant,
device_registry: dr.DeviceRegistry, device_registry: dr.DeviceRegistry,
aioclient_mock: AiohttpClientMocker, aioclient_mock: AiohttpClientMocker,
) -> None: ) -> None:
"""Verify no call is made if controller is unavailable.""" """Verify no call is made if hub is unavailable."""
clients = [ clients = [
{ {
"is_wired": False, "is_wired": False,
@ -143,12 +143,12 @@ async def test_reconnect_client_controller_unavailable(
config_entry = await setup_unifi_integration( config_entry = await setup_unifi_integration(
hass, aioclient_mock, clients_response=clients hass, aioclient_mock, clients_response=clients
) )
controller = hass.data[UNIFI_DOMAIN][config_entry.entry_id] hub = hass.data[UNIFI_DOMAIN][config_entry.entry_id]
controller.available = False hub.available = False
aioclient_mock.clear_requests() aioclient_mock.clear_requests()
aioclient_mock.post( aioclient_mock.post(
f"https://{controller.host}:1234/api/s/{controller.site}/cmd/stamgr", f"https://{hub.host}:1234/api/s/{hub.site}/cmd/stamgr",
) )
device_entry = device_registry.async_get_or_create( device_entry = device_registry.async_get_or_create(
@ -170,14 +170,14 @@ async def test_reconnect_client_unknown_mac(
device_registry: dr.DeviceRegistry, device_registry: dr.DeviceRegistry,
aioclient_mock: AiohttpClientMocker, aioclient_mock: AiohttpClientMocker,
) -> None: ) -> None:
"""Verify no call is made if trying to reconnect a mac unknown to controller.""" """Verify no call is made if trying to reconnect a mac unknown to hub."""
config_entry = await setup_unifi_integration(hass, aioclient_mock) config_entry = await setup_unifi_integration(hass, aioclient_mock)
aioclient_mock.clear_requests() aioclient_mock.clear_requests()
device_entry = device_registry.async_get_or_create( device_entry = device_registry.async_get_or_create(
config_entry_id=config_entry.entry_id, config_entry_id=config_entry.entry_id,
connections={(dr.CONNECTION_NETWORK_MAC, "mac unknown to controller")}, connections={(dr.CONNECTION_NETWORK_MAC, "mac unknown to hub")},
) )
await hass.services.async_call( await hass.services.async_call(
@ -261,11 +261,11 @@ async def test_remove_clients(
config_entry = await setup_unifi_integration( config_entry = await setup_unifi_integration(
hass, aioclient_mock, clients_all_response=clients hass, aioclient_mock, clients_all_response=clients
) )
controller = hass.data[UNIFI_DOMAIN][config_entry.entry_id] hub = hass.data[UNIFI_DOMAIN][config_entry.entry_id]
aioclient_mock.clear_requests() aioclient_mock.clear_requests()
aioclient_mock.post( aioclient_mock.post(
f"https://{controller.host}:1234/api/s/{controller.site}/cmd/stamgr", f"https://{hub.host}:1234/api/s/{hub.site}/cmd/stamgr",
) )
await hass.services.async_call(UNIFI_DOMAIN, SERVICE_REMOVE_CLIENTS, blocking=True) await hass.services.async_call(UNIFI_DOMAIN, SERVICE_REMOVE_CLIENTS, blocking=True)
@ -277,10 +277,10 @@ async def test_remove_clients(
assert await hass.config_entries.async_unload(config_entry.entry_id) assert await hass.config_entries.async_unload(config_entry.entry_id)
async def test_remove_clients_controller_unavailable( async def test_remove_clients_hub_unavailable(
hass: HomeAssistant, aioclient_mock: AiohttpClientMocker hass: HomeAssistant, aioclient_mock: AiohttpClientMocker
) -> None: ) -> None:
"""Verify no call is made if controller is unavailable.""" """Verify no call is made if UniFi Network is unavailable."""
clients = [ clients = [
{ {
"first_seen": 100, "first_seen": 100,
@ -291,8 +291,8 @@ async def test_remove_clients_controller_unavailable(
config_entry = await setup_unifi_integration( config_entry = await setup_unifi_integration(
hass, aioclient_mock, clients_all_response=clients hass, aioclient_mock, clients_all_response=clients
) )
controller = hass.data[UNIFI_DOMAIN][config_entry.entry_id] hub = hass.data[UNIFI_DOMAIN][config_entry.entry_id]
controller.available = False hub.available = False
aioclient_mock.clear_requests() aioclient_mock.clear_requests()

View file

@ -33,12 +33,7 @@ from homeassistant.helpers import entity_registry as er
from homeassistant.helpers.entity_registry import RegistryEntryDisabler from homeassistant.helpers.entity_registry import RegistryEntryDisabler
from homeassistant.util import dt as dt_util from homeassistant.util import dt as dt_util
from .test_controller import ( from .test_hub import CONTROLLER_HOST, ENTRY_CONFIG, SITE, setup_unifi_integration
CONTROLLER_HOST,
ENTRY_CONFIG,
SITE,
setup_unifi_integration,
)
from tests.common import async_fire_time_changed from tests.common import async_fire_time_changed
from tests.test_util.aiohttp import AiohttpClientMocker from tests.test_util.aiohttp import AiohttpClientMocker
@ -780,10 +775,10 @@ async def test_no_clients(
assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 0 assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 0
async def test_controller_not_client( async def test_hub_not_client(
hass: HomeAssistant, aioclient_mock: AiohttpClientMocker hass: HomeAssistant, aioclient_mock: AiohttpClientMocker
) -> None: ) -> None:
"""Test that the controller doesn't become a switch.""" """Test that the cloud key doesn't become a switch."""
await setup_unifi_integration( await setup_unifi_integration(
hass, hass,
aioclient_mock, aioclient_mock,
@ -834,7 +829,7 @@ async def test_switches(
dpigroup_response=DPI_GROUPS, dpigroup_response=DPI_GROUPS,
dpiapp_response=DPI_APPS, dpiapp_response=DPI_APPS,
) )
controller = hass.data[UNIFI_DOMAIN][config_entry.entry_id] hub = hass.data[UNIFI_DOMAIN][config_entry.entry_id]
assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 3 assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 3
@ -862,7 +857,7 @@ async def test_switches(
# Block and unblock client # Block and unblock client
aioclient_mock.clear_requests() aioclient_mock.clear_requests()
aioclient_mock.post( aioclient_mock.post(
f"https://{controller.host}:1234/api/s/{controller.site}/cmd/stamgr", f"https://{hub.host}:1234/api/s/{hub.site}/cmd/stamgr",
) )
await hass.services.async_call( await hass.services.async_call(
@ -886,7 +881,7 @@ async def test_switches(
# Enable and disable DPI # Enable and disable DPI
aioclient_mock.clear_requests() aioclient_mock.clear_requests()
aioclient_mock.put( aioclient_mock.put(
f"https://{controller.host}:1234/api/s/{controller.site}/rest/dpiapp/5f976f62e3c58f018ec7e17d", f"https://{hub.host}:1234/api/s/{hub.site}/rest/dpiapp/5f976f62e3c58f018ec7e17d",
) )
await hass.services.async_call( await hass.services.async_call(
@ -956,7 +951,7 @@ async def test_block_switches(
clients_response=[UNBLOCKED], clients_response=[UNBLOCKED],
clients_all_response=[BLOCKED], clients_all_response=[BLOCKED],
) )
controller = hass.data[UNIFI_DOMAIN][config_entry.entry_id] hub = hass.data[UNIFI_DOMAIN][config_entry.entry_id]
assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 2 assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 2
@ -986,7 +981,7 @@ async def test_block_switches(
aioclient_mock.clear_requests() aioclient_mock.clear_requests()
aioclient_mock.post( aioclient_mock.post(
f"https://{controller.host}:1234/api/s/{controller.site}/cmd/stamgr", f"https://{hub.host}:1234/api/s/{hub.site}/cmd/stamgr",
) )
await hass.services.async_call( await hass.services.async_call(
@ -1147,7 +1142,7 @@ async def test_outlet_switches(
config_entry = await setup_unifi_integration( config_entry = await setup_unifi_integration(
hass, aioclient_mock, devices_response=[test_data] hass, aioclient_mock, devices_response=[test_data]
) )
controller = hass.data[UNIFI_DOMAIN][config_entry.entry_id] hub = hass.data[UNIFI_DOMAIN][config_entry.entry_id]
assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == expected_switches assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == expected_switches
# Validate state object # Validate state object
switch_1 = hass.states.get(f"switch.{entity_id}") switch_1 = hass.states.get(f"switch.{entity_id}")
@ -1166,7 +1161,7 @@ async def test_outlet_switches(
device_id = test_data["device_id"] device_id = test_data["device_id"]
aioclient_mock.clear_requests() aioclient_mock.clear_requests()
aioclient_mock.put( aioclient_mock.put(
f"https://{controller.host}:1234/api/s/{controller.site}/rest/device/{device_id}", f"https://{hub.host}:1234/api/s/{hub.site}/rest/device/{device_id}",
) )
await hass.services.async_call( await hass.services.async_call(
@ -1338,7 +1333,7 @@ async def test_poe_port_switches(
config_entry = await setup_unifi_integration( config_entry = await setup_unifi_integration(
hass, aioclient_mock, devices_response=[DEVICE_1] hass, aioclient_mock, devices_response=[DEVICE_1]
) )
controller = hass.data[UNIFI_DOMAIN][config_entry.entry_id] hub = hass.data[UNIFI_DOMAIN][config_entry.entry_id]
assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 0 assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 0
@ -1377,7 +1372,7 @@ async def test_poe_port_switches(
# Turn off PoE # Turn off PoE
aioclient_mock.clear_requests() aioclient_mock.clear_requests()
aioclient_mock.put( aioclient_mock.put(
f"https://{controller.host}:1234/api/s/{controller.site}/rest/device/mock-id", f"https://{hub.host}:1234/api/s/{hub.site}/rest/device/mock-id",
) )
await hass.services.async_call( await hass.services.async_call(
@ -1450,7 +1445,7 @@ async def test_wlan_switches(
config_entry = await setup_unifi_integration( config_entry = await setup_unifi_integration(
hass, aioclient_mock, wlans_response=[WLAN] hass, aioclient_mock, wlans_response=[WLAN]
) )
controller = hass.data[UNIFI_DOMAIN][config_entry.entry_id] hub = hass.data[UNIFI_DOMAIN][config_entry.entry_id]
assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 1 assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 1
@ -1474,8 +1469,7 @@ async def test_wlan_switches(
# Disable WLAN # Disable WLAN
aioclient_mock.clear_requests() aioclient_mock.clear_requests()
aioclient_mock.put( aioclient_mock.put(
f"https://{controller.host}:1234/api/s/{controller.site}" f"https://{hub.host}:1234/api/s/{hub.site}" + f"/rest/wlanconf/{WLAN['_id']}",
+ f"/rest/wlanconf/{WLAN['_id']}",
) )
await hass.services.async_call( await hass.services.async_call(
@ -1531,7 +1525,7 @@ async def test_port_forwarding_switches(
config_entry = await setup_unifi_integration( config_entry = await setup_unifi_integration(
hass, aioclient_mock, port_forward_response=[_data.copy()] hass, aioclient_mock, port_forward_response=[_data.copy()]
) )
controller = hass.data[UNIFI_DOMAIN][config_entry.entry_id] hub = hass.data[UNIFI_DOMAIN][config_entry.entry_id]
assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 1 assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 1
@ -1555,7 +1549,7 @@ async def test_port_forwarding_switches(
# Disable port forward # Disable port forward
aioclient_mock.clear_requests() aioclient_mock.clear_requests()
aioclient_mock.put( aioclient_mock.put(
f"https://{controller.host}:1234/api/s/{controller.site}" f"https://{hub.host}:1234/api/s/{hub.site}"
+ f"/rest/portforward/{data['_id']}", + f"/rest/portforward/{data['_id']}",
) )

View file

@ -25,7 +25,7 @@ from homeassistant.const import (
) )
from homeassistant.core import HomeAssistant from homeassistant.core import HomeAssistant
from .test_controller import SITE, setup_unifi_integration from .test_hub import SITE, setup_unifi_integration
from tests.test_util.aiohttp import AiohttpClientMocker from tests.test_util.aiohttp import AiohttpClientMocker
@ -183,10 +183,10 @@ async def test_install(
) )
async def test_controller_state_change( async def test_hub_state_change(
hass: HomeAssistant, aioclient_mock: AiohttpClientMocker, websocket_mock hass: HomeAssistant, aioclient_mock: AiohttpClientMocker, websocket_mock
) -> None: ) -> None:
"""Verify entities state reflect on controller becoming unavailable.""" """Verify entities state reflect on hub becoming unavailable."""
await setup_unifi_integration(hass, aioclient_mock, devices_response=[DEVICE_1]) await setup_unifi_integration(hass, aioclient_mock, devices_response=[DEVICE_1])
assert len(hass.states.async_entity_ids(UPDATE_DOMAIN)) == 1 assert len(hass.states.async_entity_ids(UPDATE_DOMAIN)) == 1