From e9bbf773d63444121ef90acccdcc144226470cb6 Mon Sep 17 00:00:00 2001 From: MarkGodwin <10632972+MarkGodwin@users.noreply.github.com> Date: Mon, 30 Sep 2024 07:58:19 +0100 Subject: [PATCH] Switch oamda to use a strongly typed config entry (#127044) --- .../components/tplink_omada/__init__.py | 20 +++++++++---------- .../components/tplink_omada/binary_sensor.py | 9 ++++----- .../components/tplink_omada/device_tracker.py | 9 ++++----- .../components/tplink_omada/switch.py | 13 ++++-------- .../components/tplink_omada/update.py | 10 ++++------ 5 files changed, 26 insertions(+), 35 deletions(-) diff --git a/homeassistant/components/tplink_omada/__init__.py b/homeassistant/components/tplink_omada/__init__.py index 9945df2bbae..7890d5936fb 100644 --- a/homeassistant/components/tplink_omada/__init__.py +++ b/homeassistant/components/tplink_omada/__init__.py @@ -29,10 +29,11 @@ PLATFORMS: list[Platform] = [ ] -async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: - """Set up TP-Link Omada from a config entry.""" +type OmadaConfigEntry = ConfigEntry[OmadaSiteController] - 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): - hass.data[DOMAIN].pop(entry.entry_id) - - return unload_ok + return await hass.config_entries.async_unload_platforms(entry, PLATFORMS) 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) diff --git a/homeassistant/components/tplink_omada/binary_sensor.py b/homeassistant/components/tplink_omada/binary_sensor.py index c3941ff7595..da0c1dd9fc9 100644 --- a/homeassistant/components/tplink_omada/binary_sensor.py +++ b/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 .controller import OmadaGatewayCoordinator, OmadaSiteController +from . import OmadaConfigEntry +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: diff --git a/homeassistant/components/tplink_omada/device_tracker.py b/homeassistant/components/tplink_omada/device_tracker.py index e5a85186f24..fe78adf8847 100644 --- a/homeassistant/components/tplink_omada/device_tracker.py +++ b/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, OmadaSiteController +from .controller import OmadaClientsCoordinator _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] diff --git a/homeassistant/components/tplink_omada/switch.py b/homeassistant/components/tplink_omada/switch.py index 12d4d4039ee..26bedc5a88e 100644 --- a/homeassistant/components/tplink_omada/switch.py +++ b/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 .controller import ( - OmadaGatewayCoordinator, - OmadaSiteController, - OmadaSwitchPortCoordinator, -) +from . import OmadaConfigEntry +from .controller import OmadaGatewayCoordinator, 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 diff --git a/homeassistant/components/tplink_omada/update.py b/homeassistant/components/tplink_omada/update.py index 82c694a5ae4..d1e0a08b803 100644 --- a/homeassistant/components/tplink_omada/update.py +++ b/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 .controller import OmadaSiteController +from . import OmadaConfigEntry 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