Use config entry runtime_data in nws (home-assistant#117593)

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:

homeassistant/components/nws/diagnostics.py

@@ -4,25 +4,23 @@
 
 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"}
 
 
 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),

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(

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(

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