jeans/core
jeans/core
1
0
Fork 0
mirror of https://github.com/home-assistant/core.git synced 2025-04-25 09:47:52 +00:00
Code Issues Packages Projects Releases Wiki Activity

Break out UniFi platform registration to its own class (#112514)

Browse Source
This commit is contained in:
Robert Svensson 2024-03-09 09:18:23 +01:00 committed by GitHub
parent 9ca9d7f48f
commit f2879e6f39
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
11 changed files with 158 additions and 117 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
homeassistant/components/unifi/__init__.py
View File

@ -49,6 +49,7 @@ async def async_setup_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> b
await hass.config_entries.async_forward_entry_setups(config_entry, PLATFORMS)
hub.async_update_device_registry()
hub.entity_loader.load_entities()
if len(hass.data[UNIFI_DOMAIN]) == 1:
async_setup_services(hass)

4
homeassistant/components/unifi/button.py
View File

@ -107,9 +107,7 @@ async def async_setup_entry(
async_add_entities: AddEntitiesCallback,
) -> None:
"""Set up button platform for UniFi Network integration."""
UnifiHub.register_platform(
hass,
config_entry,
UnifiHub.get_hub(hass, config_entry).entity_loader.register_platform(
async_add_entities,
UnifiButtonEntity,
ENTITY_DESCRIPTIONS,

4
homeassistant/components/unifi/device_tracker.py
View File

@ -220,8 +220,8 @@ async def async_setup_entry(
) -> None:
"""Set up device tracker for UniFi Network integration."""
async_update_unique_id(hass, config_entry)
UnifiHub.register_platform(
hass, config_entry, async_add_entities, UnifiScannerEntity, ENTITY_DESCRIPTIONS
UnifiHub.get_hub(hass, config_entry).entity_loader.register_platform(
async_add_entities, UnifiScannerEntity, ENTITY_DESCRIPTIONS
)

6
homeassistant/components/unifi/entity.py
View File

@ -133,7 +133,7 @@ class UnifiEntity(Entity, Generic[HandlerT, ApiItemT]):
self.hub = hub
self.entity_description = description
hub.known_objects.add((description.key, obj_id))
hub.entity_loader.known_objects.add((description.key, obj_id))
self._removed = False
@ -154,7 +154,9 @@ class UnifiEntity(Entity, Generic[HandlerT, ApiItemT]):
@callback
def unregister_object() -> None:
"""Remove object ID from known_objects when unloaded."""
self.hub.known_objects.discard((description.key, self._obj_id))
self.hub.entity_loader.known_objects.discard(
(description.key, self._obj_id)
)
self.async_on_remove(unregister_object)

131
homeassistant/components/unifi/hub/entity_loader.py Normal file
View File

@ -0,0 +1,131 @@
"""UniFi Network entity loader.
Central point to load entities for the different platforms.
Make sure expected clients are available for platforms.
"""
from __future__ import annotations
from collections.abc import Iterable
from datetime import timedelta
from functools import partial
from typing import TYPE_CHECKING
from aiounifi.interfaces.api_handlers import ItemEvent
from homeassistant.core import callback
from homeassistant.helpers.dispatcher import async_dispatcher_connect
from homeassistant.helpers.entity_platform import AddEntitiesCallback
from ..entity import UnifiEntity, UnifiEntityDescription
if TYPE_CHECKING:
from .hub import UnifiHub
CHECK_HEARTBEAT_INTERVAL = timedelta(seconds=1)
class UnifiEntityLoader:
"""UniFi Network integration handling platforms for entity registration."""
def __init__(self, hub: UnifiHub) -> None:
"""Initialize the UniFi entity loader."""
self.hub = hub
self.platforms: list[
tuple[
AddEntitiesCallback,
type[UnifiEntity],
tuple[UnifiEntityDescription, ...],
bool,
]
] = []
self.known_objects: set[tuple[str, str]] = set()
"""Tuples of entity description key and object ID of loaded entities."""
@callback
def register_platform(
self,
async_add_entities: AddEntitiesCallback,
entity_class: type[UnifiEntity],
descriptions: tuple[UnifiEntityDescription, ...],
requires_admin: bool = False,
) -> None:
"""Register UniFi entity platforms."""
self.platforms.append(
(async_add_entities, entity_class, descriptions, requires_admin)
)
@callback
def load_entities(self) -> None:
"""Populate UniFi platforms with entities."""
for (
async_add_entities,
entity_class,
descriptions,
requires_admin,
) in self.platforms:
if requires_admin and not self.hub.is_admin:
continue
self._load_entities(entity_class, descriptions, async_add_entities)
@callback
def _should_add_entity(
self, description: UnifiEntityDescription, obj_id: str
) -> bool:
"""Check if entity should be added."""
return bool(
(description.key, obj_id) not in self.known_objects
and description.allowed_fn(self.hub, obj_id)
and description.supported_fn(self.hub, obj_id)
)
@callback
def _load_entities(
self,
unifi_platform_entity: type[UnifiEntity],
descriptions: tuple[UnifiEntityDescription, ...],
async_add_entities: AddEntitiesCallback,
) -> None:
"""Subscribe to UniFi API handlers and create entities."""
@callback
def async_load_entities(descriptions: Iterable[UnifiEntityDescription]) -> None:
"""Load and subscribe to UniFi endpoints."""
@callback
def _add_unifi_entities() -> None:
"""Add UniFi entity."""
async_add_entities(
unifi_platform_entity(obj_id, self.hub, description)
for description in descriptions
for obj_id in description.api_handler_fn(self.hub.api)
if self._should_add_entity(description, obj_id)
)
_add_unifi_entities()
@callback
def _create_unifi_entity(
description: UnifiEntityDescription, event: ItemEvent, obj_id: str
) -> None:
"""Create new UniFi entity on event."""
if self._should_add_entity(description, obj_id):
async_add_entities(
[unifi_platform_entity(obj_id, self.hub, description)]
)
for description in descriptions:
description.api_handler_fn(self.hub.api).subscribe(
partial(_create_unifi_entity, description), ItemEvent.ADDED
)
self.hub.config.entry.async_on_unload(
async_dispatcher_connect(
self.hub.hass,
self.hub.signal_options_update,
_add_unifi_entities,
)
)
async_load_entities(descriptions)

101
homeassistant/components/unifi/hub/hub.py
View File

@ -2,12 +2,9 @@
from __future__ import annotations
from collections.abc import Iterable
from datetime import datetime, timedelta
from functools import partial
import aiounifi
from aiounifi.interfaces.api_handlers import ItemEvent
from aiounifi.models.device import DeviceSetPoePortModeRequest
from homeassistant.config_entries import ConfigEntry
@ -19,11 +16,7 @@ from homeassistant.helpers.device_registry import (
DeviceEntryType,
DeviceInfo,
)
from homeassistant.helpers.dispatcher import (
async_dispatcher_connect,
async_dispatcher_send,
)
from homeassistant.helpers.entity_platform import AddEntitiesCallback
from homeassistant.helpers.dispatcher import async_dispatcher_send
from homeassistant.helpers.entity_registry import async_entries_for_config_entry
from homeassistant.helpers.event import async_call_later, async_track_time_interval
import homeassistant.util.dt as dt_util
@ -35,8 +28,8 @@ from ..const import (
PLATFORMS,
UNIFI_WIRELESS_CLIENTS,
)
from ..entity import UnifiEntity, UnifiEntityDescription
from .config import UnifiConfig
from .entity_loader import UnifiEntityLoader
from .websocket import UnifiWebsocket
CHECK_HEARTBEAT_INTERVAL = timedelta(seconds=1)
@ -52,6 +45,7 @@ class UnifiHub:
self.hass = hass
self.api = api
self.config = UnifiConfig.from_config_entry(config_entry)
self.entity_loader = UnifiEntityLoader(self)
self.websocket = UnifiWebsocket(hass, api, self.signal_reachable)
self.wireless_clients = hass.data[UNIFI_WIRELESS_CLIENTS]
@ -62,96 +56,21 @@ class UnifiHub:
self._cancel_heartbeat_check: CALLBACK_TYPE | None = None
self._heartbeat_time: dict[str, datetime] = {}
self.entities: dict[str, str] = {}
self.known_objects: set[tuple[str, str]] = set()
self.poe_command_queue: dict[str, dict[int, str]] = {}
self._cancel_poe_command: CALLBACK_TYPE | None = None
@callback
@staticmethod
def get_hub(hass: HomeAssistant, config_entry: ConfigEntry) -> UnifiHub:
"""Get UniFi hub from config entry."""
hub: UnifiHub = hass.data[UNIFI_DOMAIN][config_entry.entry_id]
return hub
@property
def available(self) -> bool:
"""Websocket connection state."""
return self.websocket.available
@callback
@staticmethod
def register_platform(
hass: HomeAssistant,
config_entry: ConfigEntry,
async_add_entities: AddEntitiesCallback,
entity_class: type[UnifiEntity],
descriptions: tuple[UnifiEntityDescription, ...],
requires_admin: bool = False,
) -> None:
"""Register platform for UniFi entity management."""
hub: UnifiHub = hass.data[UNIFI_DOMAIN][config_entry.entry_id]
if requires_admin and not hub.is_admin:
return
hub.register_platform_add_entities(
entity_class, descriptions, async_add_entities
)
@callback
def _async_should_add_entity(
self, description: UnifiEntityDescription, obj_id: str
) -> bool:
"""Check if entity should be added."""
return bool(
(description.key, obj_id) not in self.known_objects
and description.allowed_fn(self, obj_id)
and description.supported_fn(self, obj_id)
)
@callback
def register_platform_add_entities(
self,
unifi_platform_entity: type[UnifiEntity],
descriptions: tuple[UnifiEntityDescription, ...],
async_add_entities: AddEntitiesCallback,
) -> None:
"""Subscribe to UniFi API handlers and create entities."""
@callback
def async_load_entities(descriptions: Iterable[UnifiEntityDescription]) -> None:
"""Load and subscribe to UniFi endpoints."""
@callback
def async_add_unifi_entities() -> None:
"""Add UniFi entity."""
async_add_entities(
unifi_platform_entity(obj_id, self, description)
for description in descriptions
for obj_id in description.api_handler_fn(self.api)
if self._async_should_add_entity(description, obj_id)
)
async_add_unifi_entities()
@callback
def async_create_entity(
description: UnifiEntityDescription, event: ItemEvent, obj_id: str
) -> None:
"""Create new UniFi entity on event."""
if self._async_should_add_entity(description, obj_id):
async_add_entities(
[unifi_platform_entity(obj_id, self, description)]
)
for description in descriptions:
description.api_handler_fn(self.api).subscribe(
partial(async_create_entity, description), ItemEvent.ADDED
)
self.config.entry.async_on_unload(
async_dispatcher_connect(
self.hass,
self.signal_options_update,
async_add_unifi_entities,
)
)
async_load_entities(descriptions)
@property
def signal_reachable(self) -> str:
"""Integration specific event to signal a change in connection status."""

4
homeassistant/components/unifi/image.py
View File

@ -72,9 +72,7 @@ async def async_setup_entry(
async_add_entities: AddEntitiesCallback,
) -> None:
"""Set up image platform for UniFi Network integration."""
UnifiHub.register_platform(
hass,
config_entry,
UnifiHub.get_hub(hass, config_entry).entity_loader.register_platform(
async_add_entities,
UnifiImageEntity,
ENTITY_DESCRIPTIONS,

4
homeassistant/components/unifi/sensor.py
View File

@ -374,8 +374,8 @@ async def async_setup_entry(
async_add_entities: AddEntitiesCallback,
) -> None:
"""Set up sensors for UniFi Network integration."""
UnifiHub.register_platform(
hass, config_entry, async_add_entities, UnifiSensorEntity, ENTITY_DESCRIPTIONS
UnifiHub.get_hub(hass, config_entry).entity_loader.register_platform(
async_add_entities, UnifiSensorEntity, ENTITY_DESCRIPTIONS
)

4
homeassistant/components/unifi/switch.py
View File

@ -321,9 +321,7 @@ async def async_setup_entry(
) -> None:
"""Set up switches for UniFi Network integration."""
async_update_unique_id(hass, config_entry)
UnifiHub.register_platform(
hass,
config_entry,
UnifiHub.get_hub(hass, config_entry).entity_loader.register_platform(
async_add_entities,
UnifiSwitchEntity,
ENTITY_DESCRIPTIONS,

4
homeassistant/components/unifi/update.py
View File

@ -76,9 +76,7 @@ async def async_setup_entry(
async_add_entities: AddEntitiesCallback,
) -> None:
"""Set up update entities for UniFi Network integration."""
UnifiHub.register_platform(
hass,
config_entry,
UnifiHub.get_hub(hass, config_entry).entity_loader.register_platform(
async_add_entities,
UnifiDeviceUpdateEntity,
ENTITY_DESCRIPTIONS,

12
tests/components/unifi/test_config_flow.py
View File

@ -76,14 +76,15 @@ DEVICES = [
]
WLANS = [
{"_id": "1", "name": "SSID 1"},
{"_id": "1", "name": "SSID 1", "enabled": True},
{
"_id": "2",
"name": "SSID 2",
"name_combine_enabled": False,
"name_combine_suffix": "_IOT",
"enabled": True,
},
{"_id": "3", "name": "SSID 4", "name_combine_enabled": False},
{"_id": "3", "name": "SSID 4", "name_combine_enabled": False, "enabled": True},
]
DPI_GROUPS = [
@ -542,12 +543,7 @@ async def test_simple_option_flow(
) -> None:
"""Test simple config flow options."""
config_entry = await setup_unifi_integration(
hass,
aioclient_mock,
clients_response=CLIENTS,
wlans_response=WLANS,
dpigroup_response=DPI_GROUPS,
dpiapp_response=[],
hass, aioclient_mock, clients_response=CLIENTS
)
result = await hass.config_entries.options.async_init(