From e12b100a3743f4405f22284c3bef58daed7d8ded Mon Sep 17 00:00:00 2001 From: epenet <6771947+epenet@users.noreply.github.com> Date: Wed, 12 Feb 2025 12:49:26 +0100 Subject: [PATCH] Use runtime_data in fireservicerota (#138361) --- .../components/fireservicerota/__init__.py | 26 ++++++++----------- .../fireservicerota/binary_sensor.py | 18 ++++++------- .../components/fireservicerota/coordinator.py | 9 ++++++- .../components/fireservicerota/sensor.py | 11 +++----- .../components/fireservicerota/switch.py | 15 ++++++----- 5 files changed, 40 insertions(+), 39 deletions(-) diff --git a/homeassistant/components/fireservicerota/__init__.py b/homeassistant/components/fireservicerota/__init__.py index 6f48dcfc4bc..0f30a29cfba 100644 --- a/homeassistant/components/fireservicerota/__init__.py +++ b/homeassistant/components/fireservicerota/__init__.py @@ -4,23 +4,23 @@ from __future__ import annotations from datetime import timedelta -from homeassistant.config_entries import ConfigEntry from homeassistant.const import Platform from homeassistant.core import HomeAssistant -from .const import DATA_CLIENT, DATA_COORDINATOR, DOMAIN -from .coordinator import FireServiceRotaClient, FireServiceUpdateCoordinator +from .coordinator import ( + FireServiceConfigEntry, + FireServiceRotaClient, + FireServiceUpdateCoordinator, +) MIN_TIME_BETWEEN_UPDATES = timedelta(seconds=60) PLATFORMS = [Platform.BINARY_SENSOR, Platform.SENSOR, Platform.SWITCH] -async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: +async def async_setup_entry(hass: HomeAssistant, entry: FireServiceConfigEntry) -> bool: """Set up FireServiceRota from a config entry.""" - hass.data.setdefault(DOMAIN, {}) - client = FireServiceRotaClient(hass, entry) await client.setup() @@ -32,19 +32,15 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: await coordinator.async_config_entry_first_refresh() - hass.data[DOMAIN][entry.entry_id] = { - DATA_CLIENT: client, - DATA_COORDINATOR: coordinator, - } + entry.runtime_data = coordinator await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS) return True -async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: +async def async_unload_entry( + hass: HomeAssistant, entry: FireServiceConfigEntry +) -> bool: """Unload FireServiceRota config entry.""" - unload_ok = await hass.config_entries.async_unload_platforms(entry, PLATFORMS) - if unload_ok: - del hass.data[DOMAIN][entry.entry_id] - return unload_ok + return await hass.config_entries.async_unload_platforms(entry, PLATFORMS) diff --git a/homeassistant/components/fireservicerota/binary_sensor.py b/homeassistant/components/fireservicerota/binary_sensor.py index b8a542cf37c..be7add191c0 100644 --- a/homeassistant/components/fireservicerota/binary_sensor.py +++ b/homeassistant/components/fireservicerota/binary_sensor.py @@ -10,24 +10,22 @@ from homeassistant.core import HomeAssistant from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback from homeassistant.helpers.update_coordinator import CoordinatorEntity -from .const import DATA_CLIENT, DATA_COORDINATOR, DOMAIN as FIRESERVICEROTA_DOMAIN -from .coordinator import FireServiceRotaClient, FireServiceUpdateCoordinator +from .coordinator import ( + FireServiceConfigEntry, + FireServiceRotaClient, + FireServiceUpdateCoordinator, +) async def async_setup_entry( hass: HomeAssistant, - entry: ConfigEntry, + entry: FireServiceConfigEntry, async_add_entities: AddConfigEntryEntitiesCallback, ) -> None: """Set up FireServiceRota binary sensor based on a config entry.""" - client: FireServiceRotaClient = hass.data[FIRESERVICEROTA_DOMAIN][entry.entry_id][ - DATA_CLIENT - ] - - coordinator: FireServiceUpdateCoordinator = hass.data[FIRESERVICEROTA_DOMAIN][ - entry.entry_id - ][DATA_COORDINATOR] + coordinator = entry.runtime_data + client = coordinator.client async_add_entities([ResponseBinarySensor(coordinator, client, entry)]) diff --git a/homeassistant/components/fireservicerota/coordinator.py b/homeassistant/components/fireservicerota/coordinator.py index c452421d57b..6815bf39104 100644 --- a/homeassistant/components/fireservicerota/coordinator.py +++ b/homeassistant/components/fireservicerota/coordinator.py @@ -28,12 +28,19 @@ _LOGGER = logging.getLogger(__name__) PLATFORMS = [Platform.BINARY_SENSOR, Platform.SENSOR, Platform.SWITCH] +type FireServiceConfigEntry = ConfigEntry[FireServiceUpdateCoordinator] + class FireServiceUpdateCoordinator(DataUpdateCoordinator[dict | None]): """Data update coordinator for FireServiceRota.""" + config_entry: FireServiceConfigEntry + def __init__( - self, hass: HomeAssistant, client: FireServiceRotaClient, entry: ConfigEntry + self, + hass: HomeAssistant, + client: FireServiceRotaClient, + entry: FireServiceConfigEntry, ) -> None: """Initialize the FireServiceRota DataUpdateCoordinator.""" super().__init__( diff --git a/homeassistant/components/fireservicerota/sensor.py b/homeassistant/components/fireservicerota/sensor.py index 682c7bcc0fd..5ed65609dc8 100644 --- a/homeassistant/components/fireservicerota/sensor.py +++ b/homeassistant/components/fireservicerota/sensor.py @@ -4,27 +4,24 @@ import logging from typing import Any from homeassistant.components.sensor import SensorEntity -from homeassistant.config_entries import ConfigEntry from homeassistant.core import HomeAssistant, callback from homeassistant.helpers.dispatcher import async_dispatcher_connect from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback from homeassistant.helpers.restore_state import RestoreEntity -from .const import DATA_CLIENT, DOMAIN as FIRESERVICEROTA_DOMAIN -from .coordinator import FireServiceRotaClient +from .const import DOMAIN as FIRESERVICEROTA_DOMAIN +from .coordinator import FireServiceConfigEntry, FireServiceRotaClient _LOGGER = logging.getLogger(__name__) async def async_setup_entry( hass: HomeAssistant, - entry: ConfigEntry, + entry: FireServiceConfigEntry, async_add_entities: AddConfigEntryEntitiesCallback, ) -> None: """Set up FireServiceRota sensor based on a config entry.""" - client = hass.data[FIRESERVICEROTA_DOMAIN][entry.entry_id][DATA_CLIENT] - - async_add_entities([IncidentsSensor(client)]) + async_add_entities([IncidentsSensor(entry.runtime_data.client)]) # pylint: disable-next=hass-invalid-inheritance # needs fixing diff --git a/homeassistant/components/fireservicerota/switch.py b/homeassistant/components/fireservicerota/switch.py index 602a02a8e4a..d9fe382e4b1 100644 --- a/homeassistant/components/fireservicerota/switch.py +++ b/homeassistant/components/fireservicerota/switch.py @@ -9,21 +9,24 @@ from homeassistant.core import HomeAssistant, callback from homeassistant.helpers.dispatcher import async_dispatcher_connect from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback -from .const import DATA_CLIENT, DATA_COORDINATOR, DOMAIN as FIRESERVICEROTA_DOMAIN -from .coordinator import FireServiceRotaClient, FireServiceUpdateCoordinator +from .const import DOMAIN as FIRESERVICEROTA_DOMAIN +from .coordinator import ( + FireServiceConfigEntry, + FireServiceRotaClient, + FireServiceUpdateCoordinator, +) _LOGGER = logging.getLogger(__name__) async def async_setup_entry( hass: HomeAssistant, - entry: ConfigEntry, + entry: FireServiceConfigEntry, async_add_entities: AddConfigEntryEntitiesCallback, ) -> None: """Set up FireServiceRota switch based on a config entry.""" - client = hass.data[FIRESERVICEROTA_DOMAIN][entry.entry_id][DATA_CLIENT] - - coordinator = hass.data[FIRESERVICEROTA_DOMAIN][entry.entry_id][DATA_COORDINATOR] + coordinator = entry.runtime_data + client = coordinator.client async_add_entities([ResponseSwitch(coordinator, client, entry)])