Use runtime_data in geonetnz_quakes (#144319)

homeassistant/components/geonetnz_quakes/__init__.py

@@ -31,7 +31,6 @@ from .const import (
     DEFAULT_RADIUS,
     DEFAULT_SCAN_INTERVAL,
     DOMAIN,
-    FEED,
     PLATFORMS,
 )
 
@@ -59,6 +58,8 @@ CONFIG_SCHEMA = vol.Schema(
     extra=vol.ALLOW_EXTRA,
 )
 
+type GeonetnzQuakesConfigEntry = ConfigEntry[GeonetnzQuakesFeedEntityManager]
+
 
 async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
     """Set up the GeoNet NZ Quakes component."""
@@ -89,11 +90,10 @@ async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
     return True
 
 
-async def async_setup_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> bool:
+async def async_setup_entry(
+    hass: HomeAssistant, config_entry: GeonetnzQuakesConfigEntry
+) -> bool:
     """Set up the GeoNet NZ Quakes component as config entry."""
-    hass.data.setdefault(DOMAIN, {})
-    feeds = hass.data[DOMAIN].setdefault(FEED, {})
-
     radius = config_entry.data[CONF_RADIUS]
     if hass.config.units is US_CUSTOMARY_SYSTEM:
         radius = DistanceConverter.convert(
@@ -101,16 +101,17 @@ async def async_setup_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> b
         )
     # Create feed entity manager for all platforms.
     manager = GeonetnzQuakesFeedEntityManager(hass, config_entry, radius)
-    feeds[config_entry.entry_id] = manager
+    config_entry.runtime_data = manager
     _LOGGER.debug("Feed entity manager added for %s", config_entry.entry_id)
     await manager.async_init()
     return True
 
 
-async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+async def async_unload_entry(
+    hass: HomeAssistant, entry: GeonetnzQuakesConfigEntry
+) -> bool:
     """Unload an GeoNet NZ Quakes component config entry."""
-    manager = hass.data[DOMAIN][FEED].pop(entry.entry_id)
-    await manager.async_stop()
+    await entry.runtime_data.async_stop()
     return await hass.config_entries.async_unload_platforms(entry, PLATFORMS)
 
 

homeassistant/components/geonetnz_quakes/const.py

@@ -11,8 +11,6 @@ PLATFORMS = [Platform.GEO_LOCATION, Platform.SENSOR]
 CONF_MINIMUM_MAGNITUDE = "minimum_magnitude"
 CONF_MMI = "mmi"
 
-FEED = "feed"
-
 DEFAULT_FILTER_TIME_INTERVAL = timedelta(days=7)
 DEFAULT_MINIMUM_MAGNITUDE = 0.0
 DEFAULT_MMI = 3

homeassistant/components/geonetnz_quakes/diagnostics.py

@@ -5,28 +5,23 @@ 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_LATITUDE, CONF_LONGITUDE
 from homeassistant.core import HomeAssistant
 
-from . import GeonetnzQuakesFeedEntityManager
-from .const import DOMAIN, FEED
+from . import GeonetnzQuakesConfigEntry
 
 TO_REDACT = {CONF_LATITUDE, CONF_LONGITUDE}
 
 
 async def async_get_config_entry_diagnostics(
-    hass: HomeAssistant, config_entry: ConfigEntry
+    hass: HomeAssistant, config_entry: GeonetnzQuakesConfigEntry
 ) -> dict[str, Any]:
     """Return diagnostics for a config entry."""
     data: dict[str, Any] = {
         "info": async_redact_data(config_entry.data, TO_REDACT),
     }
 
-    manager: GeonetnzQuakesFeedEntityManager = hass.data[DOMAIN][FEED][
-        config_entry.entry_id
-    ]
-    status_info = manager.status_info()
+    status_info = config_entry.runtime_data.status_info()
     if status_info:
         data["service"] = {
             "status": status_info.status,

homeassistant/components/geonetnz_quakes/geo_location.py

@@ -9,7 +9,6 @@ from typing import Any
 from aio_geojson_geonetnz_quakes.feed_entry import GeonetnzQuakesFeedEntry
 
 from homeassistant.components.geo_location import GeolocationEvent
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import ATTR_TIME, UnitOfLength
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers import entity_registry as er
@@ -18,8 +17,7 @@ from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
 from homeassistant.util.unit_conversion import DistanceConverter
 from homeassistant.util.unit_system import US_CUSTOMARY_SYSTEM
 
-from . import GeonetnzQuakesFeedEntityManager
-from .const import DOMAIN, FEED
+from . import GeonetnzQuakesConfigEntry, GeonetnzQuakesFeedEntityManager
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -39,11 +37,11 @@ SOURCE = "geonetnz_quakes"
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    entry: ConfigEntry,
+    entry: GeonetnzQuakesConfigEntry,
     async_add_entities: AddConfigEntryEntitiesCallback,
 ) -> None:
     """Set up the GeoNet NZ Quakes Feed platform."""
-    manager: GeonetnzQuakesFeedEntityManager = hass.data[DOMAIN][FEED][entry.entry_id]
+    manager = entry.runtime_data
 
     @callback
     def async_add_geolocation(

homeassistant/components/geonetnz_quakes/sensor.py

@@ -5,13 +5,12 @@ from __future__ import annotations
 import logging
 
 from homeassistant.components.sensor import SensorEntity
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.dispatcher import async_dispatcher_connect
 from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
 from homeassistant.util import dt as dt_util
 
-from .const import DOMAIN, FEED
+from . import GeonetnzQuakesConfigEntry
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -32,11 +31,11 @@ PARALLEL_UPDATES = 0
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    entry: ConfigEntry,
+    entry: GeonetnzQuakesConfigEntry,
     async_add_entities: AddConfigEntryEntitiesCallback,
 ) -> None:
     """Set up the GeoNet NZ Quakes Feed platform."""
-    manager = hass.data[DOMAIN][FEED][entry.entry_id]
+    manager = entry.runtime_data
     sensor = GeonetnzQuakesSensor(entry.entry_id, entry.unique_id, entry.title, manager)
     async_add_entities([sensor])
     _LOGGER.debug("Sensor setup done")

tests/components/geonetnz_quakes/test_geo_location.py

@@ -5,9 +5,8 @@ from unittest.mock import patch
 
 from freezegun.api import FrozenDateTimeFactory
 
-from homeassistant.components import geonetnz_quakes
 from homeassistant.components.geo_location import ATTR_SOURCE
-from homeassistant.components.geonetnz_quakes import DEFAULT_SCAN_INTERVAL, DOMAIN, FEED
+from homeassistant.components.geonetnz_quakes import DEFAULT_SCAN_INTERVAL, DOMAIN
 from homeassistant.components.geonetnz_quakes.geo_location import (
     ATTR_DEPTH,
     ATTR_EXTERNAL_ID,
@@ -38,7 +37,7 @@ from . import _generate_mock_feed_entry
 
 from tests.common import async_fire_time_changed
 
-CONFIG = {geonetnz_quakes.DOMAIN: {CONF_RADIUS: 200}}
+CONFIG = {DOMAIN: {CONF_RADIUS: 200}}
 
 
 async def test_setup(
@@ -74,7 +73,7 @@ async def test_setup(
     freezer.move_to(utcnow)
     with patch("aio_geojson_client.feed.GeoJsonFeed.update") as mock_feed_update:
         mock_feed_update.return_value = "OK", [mock_entry_1, mock_entry_2, mock_entry_3]
-        assert await async_setup_component(hass, geonetnz_quakes.DOMAIN, CONFIG)
+        assert await async_setup_component(hass, DOMAIN, CONFIG)
         await hass.async_block_till_done()
         # Artificially trigger update and collect events.
         hass.bus.async_fire(EVENT_HOMEASSISTANT_START)
@@ -188,7 +187,7 @@ async def test_setup_imperial(
         patch("aio_geojson_client.feed.GeoJsonFeed.last_timestamp", create=True),
     ):
         mock_feed_update.return_value = "OK", [mock_entry_1]
-        assert await async_setup_component(hass, geonetnz_quakes.DOMAIN, CONFIG)
+        assert await async_setup_component(hass, DOMAIN, CONFIG)
         await hass.async_block_till_done()
         # Artificially trigger update and collect events.
         hass.bus.async_fire(EVENT_HOMEASSISTANT_START)
@@ -201,10 +200,7 @@ async def test_setup_imperial(
         )
 
         # Test conversion of 200 miles to kilometers.
-        feeds = hass.data[DOMAIN][FEED]
-        assert feeds is not None
-        assert len(feeds) == 1
-        manager = list(feeds.values())[0]
+        manager = hass.config_entries.async_loaded_entries(DOMAIN)[0].runtime_data
         # Ensure that the filter value in km is correctly set.
         assert manager._feed_manager._feed._filter_radius == 321.8688
 

tests/components/geonetnz_quakes/test_init.py

@@ -2,7 +2,6 @@
 
 from unittest.mock import patch
 
-from homeassistant.components.geonetnz_quakes import DOMAIN, FEED
 from homeassistant.core import HomeAssistant
 
 
@@ -16,8 +15,7 @@ async def test_component_unload_config_entry(hass: HomeAssistant, config_entry)
         assert await hass.config_entries.async_setup(config_entry.entry_id)
         await hass.async_block_till_done()
         assert mock_feed_manager_update.call_count == 1
-        assert hass.data[DOMAIN][FEED][config_entry.entry_id] is not None
+
         # Unload config entry.
         assert await hass.config_entries.async_unload(config_entry.entry_id)
         await hass.async_block_till_done()
-        assert hass.data[DOMAIN][FEED].get(config_entry.entry_id) is None