"""UniFi Network abstraction.""" from __future__ import annotations from datetime import datetime from typing import TYPE_CHECKING import aiounifi from homeassistant.core import Event, HomeAssistant, callback from homeassistant.helpers import device_registry as dr from homeassistant.helpers.device_registry import ( DeviceEntry, DeviceEntryType, DeviceInfo, ) from homeassistant.helpers.dispatcher import async_dispatcher_send from ..const import ATTR_MANUFACTURER, CONF_SITE_ID, DOMAIN as UNIFI_DOMAIN, PLATFORMS from .config import UnifiConfig from .entity_helper import UnifiEntityHelper from .entity_loader import UnifiEntityLoader from .websocket import UnifiWebsocket if TYPE_CHECKING: from .. import UnifiConfigEntry class UnifiHub: """Manages a single UniFi Network instance.""" def __init__( self, hass: HomeAssistant, config_entry: UnifiConfigEntry, api: aiounifi.Controller, ) -> None: """Initialize the system.""" self.hass = hass self.api = api self.config = UnifiConfig.from_config_entry(config_entry) self.entity_loader = UnifiEntityLoader(self) self._entity_helper = UnifiEntityHelper(hass, api) self.websocket = UnifiWebsocket(hass, api, self.signal_reachable) self.site = config_entry.data[CONF_SITE_ID] self.is_admin = False @property def available(self) -> bool: """Websocket connection state.""" return self.websocket.available @property def signal_heartbeat_missed(self) -> str: """Event to signal new heartbeat missed.""" return self._entity_helper.signal_heartbeat @callback def update_heartbeat(self, unique_id: str, heartbeat_expire_time: datetime) -> None: """Update device time in heartbeat monitor.""" self._entity_helper.update_heartbeat(unique_id, heartbeat_expire_time) @callback def remove_heartbeat(self, unique_id: str) -> None: """Update device time in heartbeat monitor.""" self._entity_helper.remove_heartbeat(unique_id) @callback def queue_poe_port_command( self, device_id: str, port_idx: int, poe_mode: str ) -> None: """Queue commands to execute them together per device.""" self._entity_helper.queue_poe_port_command(device_id, port_idx, poe_mode) @property def signal_reachable(self) -> str: """Integration specific event to signal a change in connection status.""" return f"unifi-reachable-{self.config.entry.entry_id}" @property def signal_options_update(self) -> str: """Event specific per UniFi entry to signal new options.""" return f"unifi-options-{self.config.entry.entry_id}" async def initialize(self) -> None: """Set up a UniFi Network instance.""" await self.entity_loader.initialize() self._entity_helper.initialize() assert self.config.entry.unique_id is not None self.is_admin = self.api.sites[self.config.entry.unique_id].role == "admin" self.config.entry.add_update_listener(self.async_config_entry_updated) @property def device_info(self) -> DeviceInfo: """UniFi Network device info.""" assert self.config.entry.unique_id is not None version: str | None = None if sysinfo := next(iter(self.api.system_information.values()), None): version = sysinfo.version return DeviceInfo( entry_type=DeviceEntryType.SERVICE, identifiers={(UNIFI_DOMAIN, self.config.entry.unique_id)}, manufacturer=ATTR_MANUFACTURER, model="UniFi Network Application", name="UniFi Network", sw_version=version, ) @callback def async_update_device_registry(self) -> DeviceEntry: """Update device registry.""" device_registry = dr.async_get(self.hass) return device_registry.async_get_or_create( config_entry_id=self.config.entry.entry_id, **self.device_info ) @staticmethod async def async_config_entry_updated( hass: HomeAssistant, config_entry: UnifiConfigEntry ) -> None: """Handle signals of config entry being updated. If config entry is updated due to reauth flow the entry might already have been reset and thus is not available. """ hub = config_entry.runtime_data hub.config = UnifiConfig.from_config_entry(config_entry) async_dispatcher_send(hass, hub.signal_options_update) @callback def shutdown(self, event: Event) -> None: """Wrap the call to unifi.close. Used as an argument to EventBus.async_listen_once. """ self.websocket.stop() async def async_reset(self) -> bool: """Reset this hub to default state. Will cancel any scheduled setup retry and will unload the config entry. """ await self.websocket.stop_and_wait() unload_ok = await self.hass.config_entries.async_unload_platforms( self.config.entry, PLATFORMS ) if not unload_ok: return False self._entity_helper.reset() return True