Strict typing of UniFi integration (#90278)

* Fix typing of UniFi controller

* Strict typing of unifi.__init__

* Strict typing of UniFi config_flow

* Strict typing of UniFi switch

* Strict typing UniFi sensor

* Strict typing UniFi device tracker

* Strict typing of UniFi

* Fix library issues related to typing

.strict-typing

@@ -311,7 +311,7 @@ homeassistant.components.trafikverket_train.*
 homeassistant.components.trafikverket_weatherstation.*
 homeassistant.components.tts.*
 homeassistant.components.twentemilieu.*
-homeassistant.components.unifi.update
+homeassistant.components.unifi.*
 homeassistant.components.unifiprotect.*
 homeassistant.components.upcloud.*
 homeassistant.components.update.*

homeassistant/components/unifi/__init__.py

@@ -64,7 +64,7 @@ async def async_setup_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> b
 
 async def async_unload_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> bool:
     """Unload a config entry."""
-    controller = hass.data[UNIFI_DOMAIN].pop(config_entry.entry_id)
+    controller: UniFiController = hass.data[UNIFI_DOMAIN].pop(config_entry.entry_id)
 
     if not hass.data[UNIFI_DOMAIN]:
         async_unload_services(hass)

homeassistant/components/unifi/controller.py

@@ -10,7 +10,7 @@ from typing import Any
 from aiohttp import CookieJar
 import aiounifi
 from aiounifi.interfaces.api_handlers import ItemEvent
-from aiounifi.websocket import WebsocketSignal, WebsocketState
+from aiounifi.websocket import WebsocketState
 import async_timeout
 
 from homeassistant.config_entries import ConfigEntry
@@ -22,7 +22,7 @@ from homeassistant.const import (
     CONF_VERIFY_SSL,
     Platform,
 )
-from homeassistant.core import HomeAssistant, callback
+from homeassistant.core import CALLBACK_TYPE, Event, HomeAssistant, callback
 from homeassistant.helpers import (
     aiohttp_client,
     device_registry as dr,
@@ -75,31 +75,32 @@ CHECK_HEARTBEAT_INTERVAL = timedelta(seconds=1)
 class UniFiController:
     """Manages a single UniFi Network instance."""
 
-    def __init__(self, hass, config_entry, api):
+    def __init__(
+        self, hass: HomeAssistant, config_entry: ConfigEntry, api: aiounifi.Controller
+    ) -> None:
         """Initialize the system."""
         self.hass = hass
         self.config_entry = config_entry
         self.api = api
 
-        api.callback = self.async_unifi_signalling_callback
+        api.ws_state_callback = self.async_unifi_ws_state_callback
 
         self.available = True
         self.wireless_clients = hass.data[UNIFI_WIRELESS_CLIENTS]
 
         self.site_id: str = ""
-        self._site_name = None
-        self._site_role = None
+        self._site_name: str | None = None
+        self._site_role: str | None = None
 
-        self._cancel_heartbeat_check = None
-        self._heartbeat_dispatch = {}
-        self._heartbeat_time = {}
+        self._cancel_heartbeat_check: CALLBACK_TYPE | None = None
+        self._heartbeat_time: dict[str, datetime] = {}
 
         self.load_config_entry_options()
 
-        self.entities = {}
+        self.entities: dict[str, str] = {}
         self.known_objects: set[tuple[str, str]] = set()
 
-    def load_config_entry_options(self):
+    def load_config_entry_options(self) -> None:
         """Store attributes to avoid property call overhead since they are called frequently."""
         options = self.config_entry.options
 
@@ -114,7 +115,7 @@ class UniFiController:
             CONF_TRACK_WIRED_CLIENTS, DEFAULT_TRACK_WIRED_CLIENTS
         )
         # Config entry option to not track devices.
-        self.option_track_devices = options.get(
+        self.option_track_devices: bool = options.get(
             CONF_TRACK_DEVICES, DEFAULT_TRACK_DEVICES
         )
         # Config entry option listing what SSIDs are being used to track clients.
@@ -133,43 +134,45 @@ class UniFiController:
         # Config entry option with list of clients to control network access.
         self.option_block_clients = options.get(CONF_BLOCK_CLIENT, [])
         # Config entry option to control DPI restriction groups.
-        self.option_dpi_restrictions = options.get(
+        self.option_dpi_restrictions: bool = options.get(
             CONF_DPI_RESTRICTIONS, DEFAULT_DPI_RESTRICTIONS
         )
 
         # Statistics sensor options
 
         # Config entry option to allow bandwidth sensors.
-        self.option_allow_bandwidth_sensors = options.get(
+        self.option_allow_bandwidth_sensors: bool = options.get(
             CONF_ALLOW_BANDWIDTH_SENSORS, DEFAULT_ALLOW_BANDWIDTH_SENSORS
         )
         # Config entry option to allow uptime sensors.
-        self.option_allow_uptime_sensors = options.get(
+        self.option_allow_uptime_sensors: bool = options.get(
             CONF_ALLOW_UPTIME_SENSORS, DEFAULT_ALLOW_UPTIME_SENSORS
         )
 
     @property
-    def host(self):
+    def host(self) -> str:
         """Return the host of this controller."""
-        return self.config_entry.data[CONF_HOST]
+        host: str = self.config_entry.data[CONF_HOST]
+        return host
 
     @property
-    def site(self):
+    def site(self) -> str:
         """Return the site of this config entry."""
-        return self.config_entry.data[CONF_SITE_ID]
+        site_id: str = self.config_entry.data[CONF_SITE_ID]
+        return site_id
 
     @property
-    def site_name(self):
+    def site_name(self) -> str | None:
         """Return the nice name of site."""
         return self._site_name
 
     @property
-    def site_role(self):
+    def site_role(self) -> str | None:
         """Return the site user role of this controller."""
         return self._site_role
 
     @property
-    def mac(self):
+    def mac(self) -> str | None:
         """Return the mac address of this controller."""
         for client in self.api.clients.values():
             if self.host == client.ip:
@@ -227,22 +230,21 @@ class UniFiController:
             async_load_entities(description)
 
     @callback
-    def async_unifi_signalling_callback(self, signal, data):
+    def async_unifi_ws_state_callback(self, state: WebsocketState) -> None:
         """Handle messages back from UniFi library."""
-        if signal == WebsocketSignal.CONNECTION_STATE:
-            if data == WebsocketState.DISCONNECTED and self.available:
-                LOGGER.warning("Lost connection to UniFi Network")
+        if state == WebsocketState.DISCONNECTED and self.available:
+            LOGGER.warning("Lost connection to UniFi Network")
 
-            if (data == WebsocketState.RUNNING and not self.available) or (
-                data == WebsocketState.DISCONNECTED and self.available
-            ):
-                self.available = data == WebsocketState.RUNNING
-                async_dispatcher_send(self.hass, self.signal_reachable)
+        if (state == WebsocketState.RUNNING and not self.available) or (
+            state == WebsocketState.DISCONNECTED and self.available
+        ):
+            self.available = state == WebsocketState.RUNNING
+            async_dispatcher_send(self.hass, self.signal_reachable)
 
-                if not self.available:
-                    self.hass.loop.call_later(RETRY_TIMER, self.reconnect, True)
-                else:
-                    LOGGER.info("Connected to UniFi Network")
+            if not self.available:
+                self.hass.loop.call_later(RETRY_TIMER, self.reconnect, True)
+            else:
+                LOGGER.info("Connected to UniFi Network")
 
     @property
     def signal_reachable(self) -> str:
@@ -259,7 +261,7 @@ class UniFiController:
         """Event specific per UniFi device tracker to signal new heartbeat missed."""
         return "unifi-heartbeat-missed"
 
-    async def initialize(self):
+    async def initialize(self) -> None:
         """Set up a UniFi Network instance."""
         await self.api.initialize()
 
@@ -291,7 +293,7 @@ class UniFiController:
                 continue
 
             client = self.api.clients_all[mac]
-            self.api.clients.process_raw([client.raw])
+            self.api.clients.process_raw([dict(client.raw)])
             LOGGER.debug(
                 "Restore disconnected client %s (%s)",
                 entry.entity_id,
@@ -319,7 +321,7 @@ class UniFiController:
             del self._heartbeat_time[unique_id]
 
     @callback
-    def _async_check_for_stale(self, *_) -> None:
+    def _async_check_for_stale(self, *_: datetime) -> None:
         """Check for any devices scheduled to be marked disconnected."""
         now = dt_util.utcnow()
 
@@ -365,7 +367,7 @@ class UniFiController:
         async_dispatcher_send(hass, controller.signal_options_update)
 
     @callback
-    def reconnect(self, log=False) -> None:
+    def reconnect(self, log: bool = False) -> None:
         """Prepare to reconnect UniFi session."""
         if log:
             LOGGER.info("Will try to reconnect to UniFi Network")
@@ -387,14 +389,14 @@ class UniFiController:
             self.hass.loop.call_later(RETRY_TIMER, self.reconnect)
 
     @callback
-    def shutdown(self, event) -> None:
+    def shutdown(self, event: Event) -> None:
         """Wrap the call to unifi.close.
 
         Used as an argument to EventBus.async_listen_once.
         """
         self.api.stop_websocket()
 
-    async def async_reset(self):
+    async def async_reset(self) -> bool:
         """Reset this controller to default state.
 
         Will cancel any scheduled setup retry and will unload
@@ -421,15 +423,15 @@ async def get_unifi_controller(
     config: MappingProxyType[str, Any],
 ) -> aiounifi.Controller:
     """Create a controller object and verify authentication."""
-    ssl_context = False
+    ssl_context: ssl.SSLContext | bool = False
 
-    if verify_ssl := bool(config.get(CONF_VERIFY_SSL)):
+    if verify_ssl := config.get(CONF_VERIFY_SSL):
         session = aiohttp_client.async_get_clientsession(hass)
         if isinstance(verify_ssl, str):
             ssl_context = ssl.create_default_context(cafile=verify_ssl)
     else:
         session = aiohttp_client.async_create_clientsession(
-            hass, verify_ssl=verify_ssl, cookie_jar=CookieJar(unsafe=True)
+            hass, verify_ssl=False, cookie_jar=CookieJar(unsafe=True)
         )
 
     controller = aiounifi.Controller(

homeassistant/components/unifi/device_tracker.py

@@ -19,7 +19,7 @@ from aiounifi.models.event import Event, EventKey
 
 from homeassistant.components.device_tracker import ScannerEntity, SourceType
 from homeassistant.config_entries import ConfigEntry
-from homeassistant.core import HomeAssistant, callback
+from homeassistant.core import Event as core_Event, HomeAssistant, callback
 from homeassistant.helpers.dispatcher import async_dispatcher_connect
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 import homeassistant.util.dt as dt_util
@@ -268,7 +268,7 @@ class UnifiScannerEntity(UnifiEntity[HandlerT, ApiItemT], ScannerEntity):
         return self._attr_unique_id
 
     @callback
-    def _make_disconnected(self, *_) -> None:
+    def _make_disconnected(self, *_: core_Event) -> None:
         """No heart beat by device."""
         self._is_connected = False
         self.async_write_ha_state()

homeassistant/components/unifi/manifest.json

@@ -8,7 +8,7 @@
   "iot_class": "local_push",
   "loggers": ["aiounifi"],
   "quality_scale": "platinum",
-  "requirements": ["aiounifi==45"],
+  "requirements": ["aiounifi==46"],
   "ssdp": [
     {
       "manufacturer": "Ubiquiti Networks",

homeassistant/components/unifi/switch.py

@@ -247,8 +247,9 @@ async def async_setup_entry(
 
     for mac in controller.option_block_clients:
         if mac not in controller.api.clients and mac in controller.api.clients_all:
-            client = controller.api.clients_all[mac]
-            controller.api.clients.process_raw([client.raw])
+            controller.api.clients.process_raw(
+                [dict(controller.api.clients_all[mac].raw)]
+            )
 
     controller.register_platform_add_entities(
         UnifiSwitchEntity, ENTITY_DESCRIPTIONS, async_add_entities

mypy.ini

@@ -2873,7 +2873,7 @@ disallow_untyped_defs = true
 warn_return_any = true
 warn_unreachable = true
 
-[mypy-homeassistant.components.unifi.update]
+[mypy-homeassistant.components.unifi.*]
 check_untyped_defs = true
 disallow_incomplete_defs = true
 disallow_subclassing_any = true

requirements_all.txt

@@ -291,7 +291,7 @@ aiosyncthing==0.5.1
 aiotractive==0.5.5
 
 # homeassistant.components.unifi
-aiounifi==45
+aiounifi==46
 
 # homeassistant.components.vlc_telnet
 aiovlc==0.1.0

requirements_test_all.txt

@@ -272,7 +272,7 @@ aiosyncthing==0.5.1
 aiotractive==0.5.5
 
 # homeassistant.components.unifi
-aiounifi==45
+aiounifi==46
 
 # homeassistant.components.vlc_telnet
 aiovlc==0.1.0