diff --git a/homeassistant/components/aurora/__init__.py b/homeassistant/components/aurora/__init__.py index cf7b48412a7..5596b82ae3f 100644 --- a/homeassistant/components/aurora/__init__.py +++ b/homeassistant/components/aurora/__init__.py @@ -1,61 +1,29 @@ """The aurora component.""" -import logging - -from auroranoaa import AuroraForecast - from homeassistant.config_entries import ConfigEntry -from homeassistant.const import CONF_LATITUDE, CONF_LONGITUDE, Platform +from homeassistant.const import Platform from homeassistant.core import HomeAssistant -from homeassistant.helpers import aiohttp_client -from .const import AURORA_API, CONF_THRESHOLD, COORDINATOR, DEFAULT_THRESHOLD, DOMAIN from .coordinator import AuroraDataUpdateCoordinator -_LOGGER = logging.getLogger(__name__) - PLATFORMS = [Platform.BINARY_SENSOR, Platform.SENSOR] +AuroraConfigEntry = ConfigEntry[AuroraDataUpdateCoordinator] -async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: + +async def async_setup_entry(hass: HomeAssistant, entry: AuroraConfigEntry) -> bool: """Set up Aurora from a config entry.""" - - conf = entry.data - options = entry.options - - session = aiohttp_client.async_get_clientsession(hass) - api = AuroraForecast(session) - - longitude = conf[CONF_LONGITUDE] - latitude = conf[CONF_LATITUDE] - threshold = options.get(CONF_THRESHOLD, DEFAULT_THRESHOLD) - - coordinator = AuroraDataUpdateCoordinator( - hass=hass, - api=api, - latitude=latitude, - longitude=longitude, - threshold=threshold, - ) + coordinator = AuroraDataUpdateCoordinator(hass=hass) await coordinator.async_config_entry_first_refresh() - hass.data.setdefault(DOMAIN, {}) - hass.data[DOMAIN][entry.entry_id] = { - COORDINATOR: coordinator, - AURORA_API: api, - } + entry.runtime_data = coordinator await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS) return True -async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: +async def async_unload_entry(hass: HomeAssistant, entry: AuroraConfigEntry) -> 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) - - return unload_ok + return await hass.config_entries.async_unload_platforms(entry, PLATFORMS) diff --git a/homeassistant/components/aurora/binary_sensor.py b/homeassistant/components/aurora/binary_sensor.py index 5c9166a0f60..f34b103e0bf 100644 --- a/homeassistant/components/aurora/binary_sensor.py +++ b/homeassistant/components/aurora/binary_sensor.py @@ -3,27 +3,28 @@ from __future__ import annotations from homeassistant.components.binary_sensor import BinarySensorEntity -from homeassistant.config_entries import ConfigEntry from homeassistant.core import HomeAssistant from homeassistant.helpers.entity_platform import AddEntitiesCallback -from .const import COORDINATOR, DOMAIN +from . import AuroraConfigEntry from .entity import AuroraEntity async def async_setup_entry( - hass: HomeAssistant, entry: ConfigEntry, async_add_entries: AddEntitiesCallback + hass: HomeAssistant, + entry: AuroraConfigEntry, + async_add_entries: AddEntitiesCallback, ) -> None: """Set up the binary_sensor platform.""" - coordinator = hass.data[DOMAIN][entry.entry_id][COORDINATOR] - - entity = AuroraSensor( - coordinator=coordinator, - translation_key="visibility_alert", + async_add_entries( + [ + AuroraSensor( + coordinator=entry.runtime_data, + translation_key="visibility_alert", + ) + ] ) - async_add_entries([entity]) - class AuroraSensor(AuroraEntity, BinarySensorEntity): """Implementation of an aurora sensor.""" diff --git a/homeassistant/components/aurora/const.py b/homeassistant/components/aurora/const.py index fef0b5e6352..7a13e85889d 100644 --- a/homeassistant/components/aurora/const.py +++ b/homeassistant/components/aurora/const.py @@ -1,8 +1,6 @@ """Constants for the Aurora integration.""" DOMAIN = "aurora" -COORDINATOR = "coordinator" -AURORA_API = "aurora_api" CONF_THRESHOLD = "forecast_threshold" DEFAULT_THRESHOLD = 75 ATTRIBUTION = "Data provided by the National Oceanic and Atmospheric Administration" diff --git a/homeassistant/components/aurora/coordinator.py b/homeassistant/components/aurora/coordinator.py index ae1101f8054..422dff83922 100644 --- a/homeassistant/components/aurora/coordinator.py +++ b/homeassistant/components/aurora/coordinator.py @@ -4,27 +4,30 @@ from __future__ import annotations from datetime import timedelta import logging +from typing import TYPE_CHECKING from aiohttp import ClientError from auroranoaa import AuroraForecast +from homeassistant.const import CONF_LATITUDE, CONF_LONGITUDE from homeassistant.core import HomeAssistant +from homeassistant.helpers.aiohttp_client import async_get_clientsession from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed +from .const import CONF_THRESHOLD, DEFAULT_THRESHOLD + +if TYPE_CHECKING: + from . import AuroraConfigEntry + _LOGGER = logging.getLogger(__name__) class AuroraDataUpdateCoordinator(DataUpdateCoordinator[int]): """Class to manage fetching data from the NOAA Aurora API.""" - def __init__( - self, - hass: HomeAssistant, - api: AuroraForecast, - latitude: float, - longitude: float, - threshold: float, - ) -> None: + config_entry: AuroraConfigEntry + + def __init__(self, hass: HomeAssistant) -> None: """Initialize the data updater.""" super().__init__( @@ -34,10 +37,12 @@ class AuroraDataUpdateCoordinator(DataUpdateCoordinator[int]): update_interval=timedelta(minutes=5), ) - self.api = api - self.latitude = int(latitude) - self.longitude = int(longitude) - self.threshold = int(threshold) + self.api = AuroraForecast(async_get_clientsession(hass)) + self.latitude = int(self.config_entry.data[CONF_LATITUDE]) + self.longitude = int(self.config_entry.data[CONF_LONGITUDE]) + self.threshold = int( + self.config_entry.options.get(CONF_THRESHOLD, DEFAULT_THRESHOLD) + ) async def _async_update_data(self) -> int: """Fetch the data from the NOAA Aurora Forecast.""" diff --git a/homeassistant/components/aurora/entity.py b/homeassistant/components/aurora/entity.py index e0dd1de3b15..317b82aed5a 100644 --- a/homeassistant/components/aurora/entity.py +++ b/homeassistant/components/aurora/entity.py @@ -1,15 +1,11 @@ """The aurora component.""" -import logging - from homeassistant.helpers.device_registry import DeviceEntryType, DeviceInfo from homeassistant.helpers.update_coordinator import CoordinatorEntity from .const import ATTRIBUTION, DOMAIN from .coordinator import AuroraDataUpdateCoordinator -_LOGGER = logging.getLogger(__name__) - class AuroraEntity(CoordinatorEntity[AuroraDataUpdateCoordinator]): """Implementation of the base Aurora Entity.""" diff --git a/homeassistant/components/aurora/sensor.py b/homeassistant/components/aurora/sensor.py index e3ae9f9cf1b..31754947843 100644 --- a/homeassistant/components/aurora/sensor.py +++ b/homeassistant/components/aurora/sensor.py @@ -3,28 +3,30 @@ from __future__ import annotations from homeassistant.components.sensor import SensorEntity, SensorStateClass -from homeassistant.config_entries import ConfigEntry from homeassistant.const import PERCENTAGE from homeassistant.core import HomeAssistant from homeassistant.helpers.entity_platform import AddEntitiesCallback -from .const import COORDINATOR, DOMAIN +from . import AuroraConfigEntry from .entity import AuroraEntity async def async_setup_entry( - hass: HomeAssistant, entry: ConfigEntry, async_add_entries: AddEntitiesCallback + hass: HomeAssistant, + entry: AuroraConfigEntry, + async_add_entries: AddEntitiesCallback, ) -> None: """Set up the sensor platform.""" - coordinator = hass.data[DOMAIN][entry.entry_id][COORDINATOR] - entity = AuroraSensor( - coordinator=coordinator, - translation_key="visibility", + async_add_entries( + [ + AuroraSensor( + coordinator=entry.runtime_data, + translation_key="visibility", + ) + ] ) - async_add_entries([entity]) - class AuroraSensor(AuroraEntity, SensorEntity): """Implementation of an aurora sensor.""" diff --git a/tests/components/aurora/test_config_flow.py b/tests/components/aurora/test_config_flow.py index a91c4eb8bc9..ada9ae9b9dd 100644 --- a/tests/components/aurora/test_config_flow.py +++ b/tests/components/aurora/test_config_flow.py @@ -56,7 +56,7 @@ async def test_form_cannot_connect(hass: HomeAssistant) -> None: ) with patch( - "homeassistant.components.aurora.AuroraForecast.get_forecast_data", + "homeassistant.components.aurora.config_flow.AuroraForecast.get_forecast_data", side_effect=ClientError, ): result = await hass.config_entries.flow.async_configure( @@ -77,7 +77,7 @@ async def test_with_unknown_error(hass: HomeAssistant) -> None: ) with patch( - "homeassistant.components.aurora.AuroraForecast.get_forecast_data", + "homeassistant.components.aurora.config_flow.AuroraForecast.get_forecast_data", side_effect=Exception, ): result = await hass.config_entries.flow.async_configure(