diff --git a/homeassistant/components/ambient_station/__init__.py b/homeassistant/components/ambient_station/__init__.py index b55a7b866cc..39586f4dbf4 100644 --- a/homeassistant/components/ambient_station/__init__.py +++ b/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(ambient.ws_disconnect(), eager_start=True) + hass.async_create_task(entry.runtime_data.ws_disconnect(), eager_start=True) return unload_ok diff --git a/homeassistant/components/ambient_station/binary_sensor.py b/homeassistant/components/ambient_station/binary_sensor.py index fc21455a00f..a79788a4c38 100644 --- a/homeassistant/components/ambient_station/binary_sensor.py +++ b/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( diff --git a/homeassistant/components/ambient_station/diagnostics.py b/homeassistant/components/ambient_station/diagnostics.py index f3508b8df38..bddbb1ab9df 100644 --- a/homeassistant/components/ambient_station/diagnostics.py +++ b/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 .const import CONF_APP_KEY, DOMAIN +from . import AmbientStationConfigEntry +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), } diff --git a/homeassistant/components/ambient_station/sensor.py b/homeassistant/components/ambient_station/sensor.py index 229ebee4fbf..dfbd2d1b4a0 100644 --- a/homeassistant/components/ambient_station/sensor.py +++ b/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 .const import ATTR_LAST_DATA, DOMAIN, TYPE_SOLARRADIATION, TYPE_SOLARRADIATION_LX +from . import AmbientStation, AmbientStationConfigEntry +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) diff --git a/tests/components/ambient_station/test_diagnostics.py b/tests/components/ambient_station/test_diagnostics.py index bc034d0e6f3..05161ba32cd 100644 --- a/tests/components/ambient_station/test_diagnostics.py +++ b/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)