From 121aa158c92e1941065be4c409b2c4e8e35082da Mon Sep 17 00:00:00 2001 From: MatthewFlamm <39341281+MatthewFlamm@users.noreply.github.com> Date: Thu, 16 May 2024 17:14:44 -0400 Subject: [PATCH] Use config entry runtime_data in nws (#117593) --- homeassistant/components/nws/__init__.py | 16 ++++++-------- homeassistant/components/nws/diagnostics.py | 10 ++++----- homeassistant/components/nws/sensor.py | 9 ++++---- homeassistant/components/nws/weather.py | 7 +++---- tests/components/nws/test_init.py | 23 ++++++--------------- 5 files changed, 23 insertions(+), 42 deletions(-) diff --git a/homeassistant/components/nws/__init__.py b/homeassistant/components/nws/__init__.py index df8cb4c329c2ac..6bcbe74a9a65ee 100644 --- a/homeassistant/components/nws/__init__.py +++ b/homeassistant/components/nws/__init__.py @@ -31,6 +31,8 @@ DEBOUNCE_TIME = 10 * 60 # in seconds +NWSConfigEntry = ConfigEntry["NWSData"] + def base_unique_id(latitude: float, longitude: float) -> str: """Return unique id for entries in configuration.""" @@ -47,7 +49,7 @@ class NWSData: coordinator_forecast_hourly: TimestampDataUpdateCoordinator[None] -async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: +async def async_setup_entry(hass: HomeAssistant, entry: NWSConfigEntry) -> bool: """Set up a National Weather Service entry.""" latitude = entry.data[CONF_LATITUDE] longitude = entry.data[CONF_LONGITUDE] @@ -130,8 +132,7 @@ def async_setup_update_forecast_hourly( hass, _LOGGER, cooldown=DEBOUNCE_TIME, immediate=True ), ) - nws_hass_data = hass.data.setdefault(DOMAIN, {}) - nws_hass_data[entry.entry_id] = NWSData( + entry.runtime_data = NWSData( nws_data, coordinator_observation, coordinator_forecast, @@ -159,14 +160,9 @@ def async_setup_update_forecast_hourly( return True -async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: +async def async_unload_entry(hass: HomeAssistant, entry: NWSConfigEntry) -> bool: """Unload a config entry.""" - unload_ok = await hass.config_entries.async_unload_platforms(entry, PLATFORMS) - if unload_ok: - hass.data[DOMAIN].pop(entry.entry_id) - if len(hass.data[DOMAIN]) == 0: - hass.data.pop(DOMAIN) - return unload_ok + return await hass.config_entries.async_unload_platforms(entry, PLATFORMS) def device_info(latitude: float, longitude: float) -> DeviceInfo: diff --git a/homeassistant/components/nws/diagnostics.py b/homeassistant/components/nws/diagnostics.py index 2ac0b2ef48853d..230991d04dfcc5 100644 --- a/homeassistant/components/nws/diagnostics.py +++ b/homeassistant/components/nws/diagnostics.py @@ -4,14 +4,12 @@ from typing import Any -from pynws import SimpleNWS - from homeassistant.components.diagnostics import async_redact_data -from homeassistant.config_entries import ConfigEntry from homeassistant.const import CONF_API_KEY, CONF_LATITUDE, CONF_LONGITUDE from homeassistant.core import HomeAssistant -from .const import CONF_STATION, DOMAIN +from . import NWSConfigEntry +from .const import CONF_STATION CONFIG_TO_REDACT = {CONF_API_KEY, CONF_LATITUDE, CONF_LONGITUDE, CONF_STATION} OBSERVATION_TO_REDACT = {"station"} @@ -19,10 +17,10 @@ async def async_get_config_entry_diagnostics( hass: HomeAssistant, - config_entry: ConfigEntry, + config_entry: NWSConfigEntry, ) -> dict[str, Any]: """Return diagnostics for a config entry.""" - nws_data: SimpleNWS = hass.data[DOMAIN][config_entry.entry_id].api + nws_data = config_entry.runtime_data.api return { "info": async_redact_data(config_entry.data, CONFIG_TO_REDACT), diff --git a/homeassistant/components/nws/sensor.py b/homeassistant/components/nws/sensor.py index 447c2dc5cf88f2..0d61e91d93b69f 100644 --- a/homeassistant/components/nws/sensor.py +++ b/homeassistant/components/nws/sensor.py @@ -12,7 +12,6 @@ SensorEntityDescription, SensorStateClass, ) -from homeassistant.config_entries import ConfigEntry from homeassistant.const import ( CONF_LATITUDE, CONF_LONGITUDE, @@ -37,8 +36,8 @@ ) from homeassistant.util.unit_system import US_CUSTOMARY_SYSTEM -from . import NWSData, base_unique_id, device_info -from .const import ATTRIBUTION, CONF_STATION, DOMAIN, OBSERVATION_VALID_TIME +from . import NWSConfigEntry, NWSData, base_unique_id, device_info +from .const import ATTRIBUTION, CONF_STATION, OBSERVATION_VALID_TIME PARALLEL_UPDATES = 0 @@ -143,10 +142,10 @@ class NWSSensorEntityDescription(SensorEntityDescription): async def async_setup_entry( - hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback + hass: HomeAssistant, entry: NWSConfigEntry, async_add_entities: AddEntitiesCallback ) -> None: """Set up the NWS weather platform.""" - nws_data: NWSData = hass.data[DOMAIN][entry.entry_id] + nws_data = entry.runtime_data station = entry.data[CONF_STATION] async_add_entities( diff --git a/homeassistant/components/nws/weather.py b/homeassistant/components/nws/weather.py index f25998f1504f20..21d9a62bbb04ef 100644 --- a/homeassistant/components/nws/weather.py +++ b/homeassistant/components/nws/weather.py @@ -23,7 +23,6 @@ Forecast, WeatherEntityFeature, ) -from homeassistant.config_entries import ConfigEntry from homeassistant.const import ( CONF_LATITUDE, CONF_LONGITUDE, @@ -38,7 +37,7 @@ from homeassistant.helpers.update_coordinator import TimestampDataUpdateCoordinator from homeassistant.util.unit_conversion import SpeedConverter, TemperatureConverter -from . import NWSData, base_unique_id, device_info +from . import NWSConfigEntry, NWSData, base_unique_id, device_info from .const import ( ATTR_FORECAST_DETAILED_DESCRIPTION, ATTRIBUTION, @@ -79,11 +78,11 @@ def convert_condition(time: str, weather: tuple[tuple[str, int | None], ...]) -> async def async_setup_entry( - hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback + hass: HomeAssistant, entry: NWSConfigEntry, async_add_entities: AddEntitiesCallback ) -> None: """Set up the NWS weather platform.""" entity_registry = er.async_get(hass) - nws_data: NWSData = hass.data[DOMAIN][entry.entry_id] + nws_data = entry.runtime_data # Remove hourly entity from legacy config entries if entity_id := entity_registry.async_get_entity_id( diff --git a/tests/components/nws/test_init.py b/tests/components/nws/test_init.py index 121da07a9ce21b..9926e530d368ca 100644 --- a/tests/components/nws/test_init.py +++ b/tests/components/nws/test_init.py @@ -1,8 +1,7 @@ """Tests for init module.""" from homeassistant.components.nws.const import DOMAIN -from homeassistant.components.weather import DOMAIN as WEATHER_DOMAIN -from homeassistant.const import STATE_UNAVAILABLE +from homeassistant.config_entries import ConfigEntryState from homeassistant.core import HomeAssistant from .const import NWS_CONFIG @@ -21,20 +20,10 @@ async def test_unload_entry(hass: HomeAssistant, mock_simple_nws) -> None: await hass.config_entries.async_setup(entry.entry_id) await hass.async_block_till_done() - assert len(hass.states.async_entity_ids(WEATHER_DOMAIN)) == 1 - assert DOMAIN in hass.data + assert len(hass.config_entries.async_entries(DOMAIN)) == 1 + assert entry.state is ConfigEntryState.LOADED - assert len(hass.data[DOMAIN]) == 1 - entries = hass.config_entries.async_entries(DOMAIN) - assert len(entries) == 1 - - assert await hass.config_entries.async_unload(entries[0].entry_id) - entities = hass.states.async_entity_ids(WEATHER_DOMAIN) - assert len(entities) == 1 - for entity in entities: - assert hass.states.get(entity).state == STATE_UNAVAILABLE - assert DOMAIN not in hass.data - - assert await hass.config_entries.async_remove(entries[0].entry_id) + assert await hass.config_entries.async_unload(entry.entry_id) await hass.async_block_till_done() - assert len(hass.states.async_entity_ids(WEATHER_DOMAIN)) == 0 + + assert entry.state is ConfigEntryState.NOT_LOADED