diff --git a/homeassistant/components/geonetnz_quakes/__init__.py b/homeassistant/components/geonetnz_quakes/__init__.py index b9443d4aed8..a1522862dca 100644 --- a/homeassistant/components/geonetnz_quakes/__init__.py +++ b/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) diff --git a/homeassistant/components/geonetnz_quakes/const.py b/homeassistant/components/geonetnz_quakes/const.py index db529a17fbe..9c0f1a08c6f 100644 --- a/homeassistant/components/geonetnz_quakes/const.py +++ b/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 diff --git a/homeassistant/components/geonetnz_quakes/diagnostics.py b/homeassistant/components/geonetnz_quakes/diagnostics.py index fbe9bf511aa..ebb6a2e9046 100644 --- a/homeassistant/components/geonetnz_quakes/diagnostics.py +++ b/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, diff --git a/homeassistant/components/geonetnz_quakes/geo_location.py b/homeassistant/components/geonetnz_quakes/geo_location.py index 96a1c3c09b2..e67d22c850f 100644 --- a/homeassistant/components/geonetnz_quakes/geo_location.py +++ b/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( diff --git a/homeassistant/components/geonetnz_quakes/sensor.py b/homeassistant/components/geonetnz_quakes/sensor.py index b8a1e2dd4db..cc4b4e16282 100644 --- a/homeassistant/components/geonetnz_quakes/sensor.py +++ b/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") diff --git a/tests/components/geonetnz_quakes/test_geo_location.py b/tests/components/geonetnz_quakes/test_geo_location.py index fd8ba81fca7..7373b207bab 100644 --- a/tests/components/geonetnz_quakes/test_geo_location.py +++ b/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 diff --git a/tests/components/geonetnz_quakes/test_init.py b/tests/components/geonetnz_quakes/test_init.py index 6730fa53ece..fd334fa57ee 100644 --- a/tests/components/geonetnz_quakes/test_init.py +++ b/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