Use runtime_data for ambient_station (#116608)

homeassistant/components/ambient_station/__init__.py

@@ -39,6 +39,8 @@ DEFAULT_SOCKET_MIN_RETRY = 15
 
 CONFIG_SCHEMA = cv.removed(DOMAIN, raise_if_present=False)
 
+AmbientStationConfigEntry = ConfigEntry["AmbientStation"]
+
 
 @callback
 def async_wm2_to_lx(value: float) -> int:
@@ -55,7 +57,9 @@ def async_hydrate_station_data(data: dict[str, Any]) -> dict[str, Any]:
     return data
 
 
-async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+async def async_setup_entry(
+    hass: HomeAssistant, entry: AmbientStationConfigEntry
+) -> bool:
     """Set up the Ambient PWS as config entry."""
     if not entry.unique_id:
         hass.config_entries.async_update_entry(
@@ -74,7 +78,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
         LOGGER.error("Config entry failed: %s", err)
         raise ConfigEntryNotReady from err
 
-    hass.data.setdefault(DOMAIN, {})[entry.entry_id] = ambient
+    entry.runtime_data = ambient
 
     async def _async_disconnect_websocket(_: Event) -> None:
         await ambient.websocket.disconnect()
@@ -88,12 +92,13 @@ 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: AmbientStationConfigEntry
+) -> bool:
     """Unload an Ambient PWS config entry."""
     unload_ok = await hass.config_entries.async_unload_platforms(entry, PLATFORMS)
     if unload_ok:
-        ambient = hass.data[DOMAIN].pop(entry.entry_id)
+        hass.async_create_task(entry.runtime_data.ws_disconnect(), eager_start=True)
-        hass.async_create_task(ambient.ws_disconnect(), eager_start=True)
 
     return unload_ok
 

homeassistant/components/ambient_station/binary_sensor.py

@@ -10,12 +10,12 @@ from homeassistant.components.binary_sensor import (
     BinarySensorEntity,
     BinarySensorEntityDescription,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import ATTR_NAME, EntityCategory
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from .const import ATTR_LAST_DATA, DOMAIN
+from . import AmbientStationConfigEntry
+from .const import ATTR_LAST_DATA
 from .entity import AmbientWeatherEntity
 
 TYPE_BATT1 = "batt1"
@@ -379,10 +379,12 @@ BINARY_SENSOR_DESCRIPTIONS = (
 
 
 async def async_setup_entry(
-    hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
+    hass: HomeAssistant,
+    entry: AmbientStationConfigEntry,
+    async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up Ambient PWS binary sensors based on a config entry."""
-    ambient = hass.data[DOMAIN][entry.entry_id]
+    ambient = entry.runtime_data
 
     async_add_entities(
         AmbientWeatherBinarySensor(

homeassistant/components/ambient_station/diagnostics.py

@@ -5,12 +5,11 @@ from __future__ import annotations
 from typing import Any
 
 from homeassistant.components.diagnostics import async_redact_data
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_API_KEY, CONF_LOCATION, CONF_UNIQUE_ID
 from homeassistant.core import HomeAssistant
 
-from . import AmbientStation
+from . import AmbientStationConfigEntry
-from .const import CONF_APP_KEY, DOMAIN
+from .const import CONF_APP_KEY
 
 CONF_API_KEY_CAMEL = "apiKey"
 CONF_APP_KEY_CAMEL = "appKey"
@@ -37,12 +36,10 @@ TO_REDACT = {
 
 
 async def async_get_config_entry_diagnostics(
-    hass: HomeAssistant, entry: ConfigEntry
+    hass: HomeAssistant, entry: AmbientStationConfigEntry
 ) -> dict[str, Any]:
     """Return diagnostics for a config entry."""
-    ambient: AmbientStation = hass.data[DOMAIN][entry.entry_id]
-
     return {
         "entry": async_redact_data(entry.as_dict(), TO_REDACT),
-        "stations": async_redact_data(ambient.stations, TO_REDACT),
+        "stations": async_redact_data(entry.runtime_data.stations, TO_REDACT),
     }

homeassistant/components/ambient_station/sensor.py

@@ -10,7 +10,6 @@ from homeassistant.components.sensor import (
     SensorEntityDescription,
     SensorStateClass,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import (
     ATTR_NAME,
     CONCENTRATION_MICROGRAMS_PER_CUBIC_METER,
@@ -30,8 +29,8 @@ from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.entity import EntityDescription
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from . import AmbientStation
+from . import AmbientStation, AmbientStationConfigEntry
-from .const import ATTR_LAST_DATA, DOMAIN, TYPE_SOLARRADIATION, TYPE_SOLARRADIATION_LX
+from .const import ATTR_LAST_DATA, TYPE_SOLARRADIATION, TYPE_SOLARRADIATION_LX
 from .entity import AmbientWeatherEntity
 
 TYPE_24HOURRAININ = "24hourrainin"
@@ -661,10 +660,12 @@ SENSOR_DESCRIPTIONS = (
 
 
 async def async_setup_entry(
-    hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
+    hass: HomeAssistant,
+    entry: AmbientStationConfigEntry,
+    async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up Ambient PWS sensors based on a config entry."""
-    ambient = hass.data[DOMAIN][entry.entry_id]
+    ambient = entry.runtime_data
 
     async_add_entities(
         AmbientWeatherSensor(ambient, mac_address, station[ATTR_NAME], description)

tests/components/ambient_station/test_diagnostics.py

@@ -2,7 +2,7 @@
 
 from syrupy import SnapshotAssertion
 
-from homeassistant.components.ambient_station import DOMAIN
+from homeassistant.components.ambient_station import AmbientStationConfigEntry
 from homeassistant.core import HomeAssistant
 
 from tests.components.diagnostics import get_diagnostics_for_config_entry
@@ -11,14 +11,14 @@ from tests.typing import ClientSessionGenerator
 
 async def test_entry_diagnostics(
     hass: HomeAssistant,
-    config_entry,
+    config_entry: AmbientStationConfigEntry,
     hass_client: ClientSessionGenerator,
     data_station,
     setup_config_entry,
     snapshot: SnapshotAssertion,
 ) -> None:
     """Test config entry diagnostics."""
-    ambient = hass.data[DOMAIN][config_entry.entry_id]
+    ambient = config_entry.runtime_data
     ambient.stations = data_station
     assert (
         await get_diagnostics_for_config_entry(hass, hass_client, config_entry)