diff --git a/homeassistant/components/launch_library/__init__.py b/homeassistant/components/launch_library/__init__.py index 6bfd3bc9adf..9b29af194e7 100644 --- a/homeassistant/components/launch_library/__init__.py +++ b/homeassistant/components/launch_library/__init__.py @@ -2,61 +2,20 @@ from __future__ import annotations -from datetime import timedelta -import logging -from typing import TypedDict - -from pylaunches import PyLaunches, PyLaunchesError -from pylaunches.types import Launch, StarshipResponse - from homeassistant.config_entries import ConfigEntry from homeassistant.const import Platform 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 DOMAIN - -_LOGGER = logging.getLogger(__name__) +from .coordinator import LaunchLibraryCoordinator PLATFORMS = [Platform.SENSOR] -class LaunchLibraryData(TypedDict): - """Typed dict representation of data returned from pylaunches.""" - - upcoming_launches: list[Launch] - starship_events: StarshipResponse - - async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: """Set up this integration using UI.""" - hass.data.setdefault(DOMAIN, {}) - - session = async_get_clientsession(hass) - launches = PyLaunches(session) - - async def async_update() -> LaunchLibraryData: - try: - return LaunchLibraryData( - upcoming_launches=await launches.launch_upcoming( - filters={"limit": 1, "hide_recent_previous": "True"}, - ), - starship_events=await launches.dashboard_starship(), - ) - except PyLaunchesError as ex: - raise UpdateFailed(ex) from ex - - coordinator = DataUpdateCoordinator( - hass, - _LOGGER, - config_entry=entry, - name=DOMAIN, - update_method=async_update, - update_interval=timedelta(hours=1), - ) - + coordinator = LaunchLibraryCoordinator(hass, entry) await coordinator.async_config_entry_first_refresh() hass.data[DOMAIN] = coordinator diff --git a/homeassistant/components/launch_library/coordinator.py b/homeassistant/components/launch_library/coordinator.py new file mode 100644 index 00000000000..b88bc105630 --- /dev/null +++ b/homeassistant/components/launch_library/coordinator.py @@ -0,0 +1,60 @@ +"""DataUpdateCoordinator for the launch_library integration.""" + +from __future__ import annotations + +from datetime import timedelta +import logging +from typing import TypedDict + +from pylaunches import PyLaunches, PyLaunchesError +from pylaunches.types import Launch, StarshipResponse + +from homeassistant.config_entries import ConfigEntry +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 DOMAIN + +_LOGGER = logging.getLogger(__name__) + + +class LaunchLibraryData(TypedDict): + """Typed dict representation of data returned from pylaunches.""" + + upcoming_launches: list[Launch] + starship_events: StarshipResponse + + +class LaunchLibraryCoordinator(DataUpdateCoordinator[LaunchLibraryData]): + """Class to manage fetching Launch Library data.""" + + config_entry: ConfigEntry + + def __init__( + self, + hass: HomeAssistant, + entry: ConfigEntry, + ) -> None: + """Initialize the coordinator.""" + super().__init__( + hass, + _LOGGER, + config_entry=entry, + name=DOMAIN, + update_interval=timedelta(hours=1), + ) + session = async_get_clientsession(hass) + self._launches = PyLaunches(session) + + async def _async_update_data(self) -> LaunchLibraryData: + """Fetch data from Launch Library.""" + try: + return LaunchLibraryData( + upcoming_launches=await self._launches.launch_upcoming( + filters={"limit": 1, "hide_recent_previous": "True"}, + ), + starship_events=await self._launches.dashboard_starship(), + ) + except PyLaunchesError as ex: + raise UpdateFailed(ex) from ex diff --git a/homeassistant/components/launch_library/diagnostics.py b/homeassistant/components/launch_library/diagnostics.py index 75541598ef5..d96d5fed7f5 100644 --- a/homeassistant/components/launch_library/diagnostics.py +++ b/homeassistant/components/launch_library/diagnostics.py @@ -8,10 +8,9 @@ from pylaunches.types import Event, Launch from homeassistant.config_entries import ConfigEntry from homeassistant.core import HomeAssistant -from homeassistant.helpers.update_coordinator import DataUpdateCoordinator -from . import LaunchLibraryData from .const import DOMAIN +from .coordinator import LaunchLibraryCoordinator async def async_get_config_entry_diagnostics( @@ -20,7 +19,7 @@ async def async_get_config_entry_diagnostics( ) -> dict[str, Any]: """Return diagnostics for a config entry.""" - coordinator: DataUpdateCoordinator[LaunchLibraryData] = hass.data[DOMAIN] + coordinator: LaunchLibraryCoordinator = hass.data[DOMAIN] if coordinator.data is None: return {} diff --git a/homeassistant/components/launch_library/sensor.py b/homeassistant/components/launch_library/sensor.py index 201b4c8f037..e844744c834 100644 --- a/homeassistant/components/launch_library/sensor.py +++ b/homeassistant/components/launch_library/sensor.py @@ -19,14 +19,11 @@ from homeassistant.const import CONF_NAME, PERCENTAGE from homeassistant.core import HomeAssistant, callback from homeassistant.helpers.device_registry import DeviceEntryType, DeviceInfo from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback -from homeassistant.helpers.update_coordinator import ( - CoordinatorEntity, - DataUpdateCoordinator, -) +from homeassistant.helpers.update_coordinator import CoordinatorEntity from homeassistant.util.dt import parse_datetime -from . import LaunchLibraryData from .const import DOMAIN +from .coordinator import LaunchLibraryCoordinator DEFAULT_NEXT_LAUNCH_NAME = "Next launch" @@ -126,7 +123,7 @@ async def async_setup_entry( ) -> None: """Set up the sensor platform.""" name = entry.data.get(CONF_NAME, DEFAULT_NEXT_LAUNCH_NAME) - coordinator: DataUpdateCoordinator[LaunchLibraryData] = hass.data[DOMAIN] + coordinator: LaunchLibraryCoordinator = hass.data[DOMAIN] async_add_entities( LaunchLibrarySensor( @@ -139,9 +136,7 @@ async def async_setup_entry( ) -class LaunchLibrarySensor( - CoordinatorEntity[DataUpdateCoordinator[LaunchLibraryData]], SensorEntity -): +class LaunchLibrarySensor(CoordinatorEntity[LaunchLibraryCoordinator], SensorEntity): """Representation of the next launch sensors.""" _attr_attribution = "Data provided by Launch Library." @@ -151,7 +146,7 @@ class LaunchLibrarySensor( def __init__( self, - coordinator: DataUpdateCoordinator[LaunchLibraryData], + coordinator: LaunchLibraryCoordinator, entry_id: str, description: LaunchLibrarySensorEntityDescription, name: str,