Use runtime_data in weatherflow_cloud integration (#168624)

Co-authored-by: Michael <35783820+mib1185@users.noreply.github.com>

homeassistant/components/weatherflow_cloud/__init__.py

@@ -3,19 +3,19 @@
 from __future__ import annotations
 
 import asyncio
-from dataclasses import dataclass
 
 from weatherflow4py.api import WeatherFlowRestAPI
 from weatherflow4py.ws import WeatherFlowWebsocketAPI
 
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_API_TOKEN, Platform
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.aiohttp_client import async_get_clientsession
 
-from .const import DOMAIN, LOGGER
+from .const import LOGGER
 from .coordinator import (
+    WeatherFlowCloudConfigEntry,
     WeatherFlowCloudUpdateCoordinatorREST,
+    WeatherFlowCoordinators,
     WeatherFlowObservationCoordinator,
     WeatherFlowWindCoordinator,
 )
@@ -23,16 +23,9 @@ from .coordinator import (
 PLATFORMS: list[Platform] = [Platform.SENSOR, Platform.WEATHER]
 
 
-@dataclass
-class WeatherFlowCoordinators:
-    """Data Class for Entry Data."""
-
-    rest: WeatherFlowCloudUpdateCoordinatorREST
-    wind: WeatherFlowWindCoordinator
-    observation: WeatherFlowObservationCoordinator
-
-
-async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+async def async_setup_entry(
+    hass: HomeAssistant, entry: WeatherFlowCloudConfigEntry
+) -> bool:
     """Set up WeatherFlowCloud from a config entry."""
 
     LOGGER.debug("Initializing WeatherFlowCloudDataUpdateCoordinatorREST coordinator")
@@ -82,7 +75,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
         websocket_observation_coordinator.async_setup(),
     )
 
-    hass.data.setdefault(DOMAIN, {})[entry.entry_id] = WeatherFlowCoordinators(
+    entry.runtime_data = WeatherFlowCoordinators(
         rest_data_coordinator,
         websocket_wind_coordinator,
         websocket_observation_coordinator,
@@ -100,10 +93,8 @@ 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: WeatherFlowCloudConfigEntry
+) -> 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)

homeassistant/components/weatherflow_cloud/coordinator.py

@@ -1,6 +1,9 @@
 """Improved coordinator design with better type safety."""
 
+from __future__ import annotations
+
 from abc import ABC, abstractmethod
+from dataclasses import dataclass
 from datetime import timedelta
 
 from aiohttp import ClientResponseError
@@ -29,13 +32,27 @@ from homeassistant.util.ssl import client_context
 from .const import DOMAIN, LOGGER
 
 
+@dataclass
+class WeatherFlowCoordinators:
+    """Data Class for Entry Data."""
+
+    rest: WeatherFlowCloudUpdateCoordinatorREST
+    wind: WeatherFlowWindCoordinator
+    observation: WeatherFlowObservationCoordinator
+
+
+type WeatherFlowCloudConfigEntry = ConfigEntry[WeatherFlowCoordinators]
+
+
 class BaseWeatherFlowCoordinator[T](DataUpdateCoordinator[dict[int, T]], ABC):
     """Base class for WeatherFlow coordinators."""
 
+    config_entry: WeatherFlowCloudConfigEntry
+
     def __init__(
         self,
         hass: HomeAssistant,
-        config_entry: ConfigEntry,
+        config_entry: WeatherFlowCloudConfigEntry,
         rest_api: WeatherFlowRestAPI,
         stations: StationsResponseREST,
         update_interval: timedelta | None = None,
@@ -70,7 +87,7 @@ class WeatherFlowCloudUpdateCoordinatorREST(
     def __init__(
         self,
         hass: HomeAssistant,
-        config_entry: ConfigEntry,
+        config_entry: WeatherFlowCloudConfigEntry,
         rest_api: WeatherFlowRestAPI,
         stations: StationsResponseREST,
     ) -> None:
@@ -111,7 +128,7 @@ class BaseWebsocketCoordinator[T](BaseWeatherFlowCoordinator[dict[int, T | None]
     def __init__(
         self,
         hass: HomeAssistant,
-        config_entry: ConfigEntry,
+        config_entry: WeatherFlowCloudConfigEntry,
         rest_api: WeatherFlowRestAPI,
         websocket_api: WeatherFlowWebsocketAPI,
         stations: StationsResponseREST,

homeassistant/components/weatherflow_cloud/sensor.py

@@ -20,7 +20,6 @@ from homeassistant.components.sensor import (
     SensorEntityDescription,
     SensorStateClass,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import (
     EntityCategory,
     UnitOfLength,
@@ -34,9 +33,12 @@ from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
 from homeassistant.helpers.typing import StateType
 from homeassistant.util.dt import UTC
 
-from . import WeatherFlowCloudUpdateCoordinatorREST, WeatherFlowCoordinators
-from .const import DOMAIN
-from .coordinator import WeatherFlowObservationCoordinator, WeatherFlowWindCoordinator
+from .coordinator import (
+    WeatherFlowCloudConfigEntry,
+    WeatherFlowCloudUpdateCoordinatorREST,
+    WeatherFlowObservationCoordinator,
+    WeatherFlowWindCoordinator,
+)
 from .entity import WeatherFlowCloudEntity
 
 PRECIPITATION_TYPE = {
@@ -350,15 +352,15 @@ WF_SENSORS: tuple[WeatherFlowCloudSensorEntityDescription, ...] = (
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    entry: ConfigEntry,
+    entry: WeatherFlowCloudConfigEntry,
     async_add_entities: AddConfigEntryEntitiesCallback,
 ) -> None:
     """Set up WeatherFlow sensors based on a config entry."""
 
-    coordinators: WeatherFlowCoordinators = hass.data[DOMAIN][entry.entry_id]
+    coordinators = entry.runtime_data
     rest_coordinator = coordinators.rest
-    wind_coordinator = coordinators.wind  # Now properly typed
-    observation_coordinator = coordinators.observation  # Now properly typed
+    wind_coordinator = coordinators.wind
+    observation_coordinator = coordinators.observation
 
     entities: list[SensorEntity] = [
         WeatherFlowCloudSensorREST(rest_coordinator, sensor_description, station_id)

homeassistant/components/weatherflow_cloud/weather.py

@@ -9,7 +9,6 @@ from homeassistant.components.weather import (
     SingleCoordinatorWeatherEntity,
     WeatherEntityFeature,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import (
     UnitOfPrecipitationDepth,
     UnitOfPressure,
@@ -19,18 +18,21 @@ from homeassistant.const import (
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
 
-from . import WeatherFlowCloudUpdateCoordinatorREST, WeatherFlowCoordinators
-from .const import DOMAIN, STATE_MAP
+from .const import STATE_MAP
+from .coordinator import (
+    WeatherFlowCloudConfigEntry,
+    WeatherFlowCloudUpdateCoordinatorREST,
+)
 from .entity import WeatherFlowCloudEntity
 
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    config_entry: ConfigEntry,
+    config_entry: WeatherFlowCloudConfigEntry,
     async_add_entities: AddConfigEntryEntitiesCallback,
 ) -> None:
     """Add a weather entity from a config_entry."""
-    coordinators: WeatherFlowCoordinators = hass.data[DOMAIN][config_entry.entry_id]
+    coordinators = config_entry.runtime_data
 
     async_add_entities(
         [

tests/components/weatherflow_cloud/test_sensor.py

@@ -8,7 +8,7 @@ import pytest
 from syrupy.assertion import SnapshotAssertion
 from weatherflow4py.models.rest.observation import ObservationStationREST
 
-from homeassistant.components.weatherflow_cloud import DOMAIN
+from homeassistant.components.weatherflow_cloud.const import DOMAIN
 from homeassistant.components.weatherflow_cloud.coordinator import (
     WeatherFlowObservationCoordinator,
     WeatherFlowWindCoordinator,