Switch oamda to use a strongly typed config entry (#127044)

homeassistant/components/tplink_omada/__init__.py

@@ -29,10 +29,11 @@ PLATFORMS: list[Platform] = [
 ]
 
 
-async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+type OmadaConfigEntry = ConfigEntry[OmadaSiteController]
-    """Set up TP-Link Omada from a config entry."""
 
-    hass.data.setdefault(DOMAIN, {})
+
+async def async_setup_entry(hass: HomeAssistant, entry: OmadaConfigEntry) -> bool:
+    """Set up TP-Link Omada from a config entry."""
 
     try:
         client = await create_omada_client(hass, entry.data)
@@ -56,7 +57,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     controller = OmadaSiteController(hass, site_client)
     await controller.initialize_first_refresh()
 
-    hass.data[DOMAIN][entry.entry_id] = controller
+    entry.runtime_data = controller
 
     _remove_old_devices(hass, entry, controller.devices_coordinator.data)
 
@@ -65,16 +66,15 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     return True
 
 
-async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+async def async_unload_entry(hass: HomeAssistant, entry: OmadaConfigEntry) -> bool:
     """Unload a config entry."""
-    if unload_ok := await hass.config_entries.async_unload_platforms(entry, PLATFORMS):
+    return await hass.config_entries.async_unload_platforms(entry, PLATFORMS)
-        hass.data[DOMAIN].pop(entry.entry_id)
-
-    return unload_ok
 
 
 def _remove_old_devices(
-    hass: HomeAssistant, entry: ConfigEntry, omada_devices: dict[str, OmadaListDevice]
+    hass: HomeAssistant,
+    entry: OmadaConfigEntry,
+    omada_devices: dict[str, OmadaListDevice],
 ) -> None:
     device_registry = dr.async_get(hass)
 

homeassistant/components/tplink_omada/binary_sensor.py

@@ -17,22 +17,21 @@ from homeassistant.components.binary_sensor import (
     BinarySensorEntity,
     BinarySensorEntityDescription,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from .const import DOMAIN
+from . import OmadaConfigEntry
-from .controller import OmadaGatewayCoordinator, OmadaSiteController
+from .controller import OmadaGatewayCoordinator
 from .entity import OmadaDeviceEntity
 
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    config_entry: ConfigEntry,
+    config_entry: OmadaConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up binary sensors."""
-    controller: OmadaSiteController = hass.data[DOMAIN][config_entry.entry_id]
+    controller = config_entry.runtime_data
 
     gateway_coordinator = controller.gateway_coordinator
     if not gateway_coordinator:

homeassistant/components/tplink_omada/device_tracker.py

@@ -5,26 +5,25 @@ import logging
 from tplink_omada_client.clients import OmadaWirelessClient
 
 from homeassistant.components.device_tracker import ScannerEntity
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.update_coordinator import CoordinatorEntity
 
+from . import OmadaConfigEntry
 from .config_flow import CONF_SITE
-from .const import DOMAIN
+from .controller import OmadaClientsCoordinator
-from .controller import OmadaClientsCoordinator, OmadaSiteController
 
 _LOGGER = logging.getLogger(__name__)
 
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    config_entry: ConfigEntry,
+    config_entry: OmadaConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up device trackers and scanners."""
 
-    controller: OmadaSiteController = hass.data[DOMAIN][config_entry.entry_id]
+    controller = config_entry.runtime_data
 
     site_id = config_entry.data[CONF_SITE]
 

homeassistant/components/tplink_omada/switch.py

@@ -20,17 +20,12 @@ from tplink_omada_client.devices import (
 from tplink_omada_client.omadasiteclient import GatewayPortSettings
 
 from homeassistant.components.switch import SwitchEntity, SwitchEntityDescription
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import EntityCategory
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from .const import DOMAIN
+from . import OmadaConfigEntry
-from .controller import (
+from .controller import OmadaGatewayCoordinator, OmadaSwitchPortCoordinator
-    OmadaGatewayCoordinator,
-    OmadaSiteController,
-    OmadaSwitchPortCoordinator,
-)
 from .coordinator import OmadaCoordinator
 from .entity import OmadaDeviceEntity
 
@@ -41,11 +36,11 @@ TCoordinator = TypeVar("TCoordinator", bound="OmadaCoordinator[Any]")
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    config_entry: ConfigEntry,
+    config_entry: OmadaConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up switches."""
-    controller: OmadaSiteController = hass.data[DOMAIN][config_entry.entry_id]
+    controller = config_entry.runtime_data
     omada_client = controller.omada_client
 
     # Naming fun. Omada switches, as in the network hardware

homeassistant/components/tplink_omada/update.py

@@ -14,13 +14,11 @@ from homeassistant.components.update import (
     UpdateEntity,
     UpdateEntityFeature,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.exceptions import HomeAssistantError
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from .const import DOMAIN
+from . import OmadaConfigEntry
-from .controller import OmadaSiteController
 from .coordinator import POLL_DEVICES, OmadaCoordinator, OmadaDevicesCoordinator
 from .entity import OmadaDeviceEntity
 
@@ -40,7 +38,7 @@ class OmadaFirmwareUpdateCoordinator(OmadaCoordinator[FirmwareUpdateStatus]):  #
     def __init__(
         self,
         hass: HomeAssistant,
-        config_entry: ConfigEntry,
+        config_entry: OmadaConfigEntry,
         omada_client: OmadaSiteClient,
         devices_coordinator: OmadaDevicesCoordinator,
     ) -> None:
@@ -92,11 +90,11 @@ class OmadaFirmwareUpdateCoordinator(OmadaCoordinator[FirmwareUpdateStatus]):  #
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    config_entry: ConfigEntry,
+    config_entry: OmadaConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up switches."""
-    controller: OmadaSiteController = hass.data[DOMAIN][config_entry.entry_id]
+    controller = config_entry.runtime_data
 
     devices = controller.devices_coordinator.data