Clean up aurora (#117315)

* Clean up aurora

* Fix

* Fix

* Fix

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."""
-
+    coordinator = AuroraDataUpdateCoordinator(hass=hass)
-    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,
-    )
 
     await coordinator.async_config_entry_first_refresh()
 
-    hass.data.setdefault(DOMAIN, {})
+    entry.runtime_data = coordinator
-    hass.data[DOMAIN][entry.entry_id] = {
-        COORDINATOR: coordinator,
-        AURORA_API: api,
-    }
 
     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)
+    return await hass.config_entries.async_unload_platforms(entry, PLATFORMS)
-
-    if unload_ok:
-        hass.data[DOMAIN].pop(entry.entry_id)
-
-    return unload_ok

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]
+    async_add_entries(
-
+        [
-    entity = AuroraSensor(
+            AuroraSensor(
-        coordinator=coordinator,
+                coordinator=entry.runtime_data,
-        translation_key="visibility_alert",
+                translation_key="visibility_alert",
+            )
+        ]
     )
 
-    async_add_entries([entity])
-
 
 class AuroraSensor(AuroraEntity, BinarySensorEntity):
     """Implementation of an aurora sensor."""

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"

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__(
+    config_entry: AuroraConfigEntry
-        self,
+
-        hass: HomeAssistant,
+    def __init__(self, hass: HomeAssistant) -> None:
-        api: AuroraForecast,
-        latitude: float,
-        longitude: float,
-        threshold: float,
-    ) -> None:
         """Initialize the data updater."""
 
         super().__init__(
@@ -34,10 +37,12 @@ class AuroraDataUpdateCoordinator(DataUpdateCoordinator[int]):
             update_interval=timedelta(minutes=5),
         )
 
-        self.api = api
+        self.api = AuroraForecast(async_get_clientsession(hass))
-        self.latitude = int(latitude)
+        self.latitude = int(self.config_entry.data[CONF_LATITUDE])
-        self.longitude = int(longitude)
+        self.longitude = int(self.config_entry.data[CONF_LONGITUDE])
-        self.threshold = int(threshold)
+        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."""

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."""

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(
+    async_add_entries(
-        coordinator=coordinator,
+        [
-        translation_key="visibility",
+            AuroraSensor(
+                coordinator=entry.runtime_data,
+                translation_key="visibility",
+            )
+        ]
     )
 
-    async_add_entries([entity])
-
 
 class AuroraSensor(AuroraEntity, SensorEntity):
     """Implementation of an aurora sensor."""

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(