Use config entry runtime data in Trafikverket Train (#116556)

This commit is contained in:
G Johansson 2024-05-01 22:08:19 +02:00 committed by GitHub
parent b089f89f14
commit 3bf67f3ddd
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 24 additions and 17 deletions

View File

@ -16,11 +16,13 @@ from homeassistant.exceptions import ConfigEntryAuthFailed, ConfigEntryNotReady
from homeassistant.helpers import entity_registry as er
from homeassistant.helpers.aiohttp_client import async_get_clientsession
from .const import CONF_FILTER_PRODUCT, CONF_FROM, CONF_TO, DOMAIN, PLATFORMS
from .const import CONF_FROM, CONF_TO, PLATFORMS
from .coordinator import TVDataUpdateCoordinator
TVTrainConfigEntry = ConfigEntry[TVDataUpdateCoordinator]
async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
async def async_setup_entry(hass: HomeAssistant, entry: TVTrainConfigEntry) -> bool:
"""Set up Trafikverket Train from a config entry."""
http_session = async_get_clientsession(hass)
@ -37,11 +39,9 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
f" {entry.data[CONF_TO]}. Error: {error} "
) from error
coordinator = TVDataUpdateCoordinator(
hass, entry, to_station, from_station, entry.options.get(CONF_FILTER_PRODUCT)
)
coordinator = TVDataUpdateCoordinator(hass, to_station, from_station)
await coordinator.async_config_entry_first_refresh()
hass.data.setdefault(DOMAIN, {})[entry.entry_id] = coordinator
entry.runtime_data = coordinator
entity_reg = er.async_get(hass)
entries = er.async_entries_for_config_entry(entity_reg, entry.entry_id)

View File

@ -5,6 +5,7 @@ from __future__ import annotations
from dataclasses import dataclass
from datetime import datetime, time, timedelta
import logging
from typing import TYPE_CHECKING
from pytrafikverket import TrafikverketTrain
from pytrafikverket.exceptions import (
@ -14,7 +15,6 @@ from pytrafikverket.exceptions import (
)
from pytrafikverket.trafikverket_train import StationInfo, TrainStop
from homeassistant.config_entries import ConfigEntry
from homeassistant.const import CONF_API_KEY, CONF_WEEKDAY
from homeassistant.core import HomeAssistant
from homeassistant.exceptions import ConfigEntryAuthFailed
@ -22,9 +22,12 @@ from homeassistant.helpers.aiohttp_client import async_get_clientsession
from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed
from homeassistant.util import dt as dt_util
from .const import CONF_TIME, DOMAIN
from .const import CONF_FILTER_PRODUCT, CONF_TIME, DOMAIN
from .util import next_departuredate
if TYPE_CHECKING:
from . import TVTrainConfigEntry
@dataclass
class TrainData:
@ -65,13 +68,13 @@ def _get_as_joined(information: list[str] | None) -> str | None:
class TVDataUpdateCoordinator(DataUpdateCoordinator[TrainData]):
"""A Trafikverket Data Update Coordinator."""
config_entry: TVTrainConfigEntry
def __init__(
self,
hass: HomeAssistant,
entry: ConfigEntry,
to_station: StationInfo,
from_station: StationInfo,
filter_product: str | None,
) -> None:
"""Initialize the Trafikverket coordinator."""
super().__init__(
@ -81,13 +84,15 @@ class TVDataUpdateCoordinator(DataUpdateCoordinator[TrainData]):
update_interval=TIME_BETWEEN_UPDATES,
)
self._train_api = TrafikverketTrain(
async_get_clientsession(hass), entry.data[CONF_API_KEY]
async_get_clientsession(hass), self.config_entry.data[CONF_API_KEY]
)
self.from_station: StationInfo = from_station
self.to_station: StationInfo = to_station
self._time: time | None = dt_util.parse_time(entry.data[CONF_TIME])
self._weekdays: list[str] = entry.data[CONF_WEEKDAY]
self._filter_product = filter_product
self._time: time | None = dt_util.parse_time(self.config_entry.data[CONF_TIME])
self._weekdays: list[str] = self.config_entry.data[CONF_WEEKDAY]
self._filter_product: str | None = self.config_entry.options.get(
CONF_FILTER_PRODUCT
)
async def _async_update_data(self) -> TrainData:
"""Fetch data from Trafikverket."""

View File

@ -12,7 +12,6 @@ from homeassistant.components.sensor import (
SensorEntity,
SensorEntityDescription,
)
from homeassistant.config_entries import ConfigEntry
from homeassistant.const import CONF_NAME, UnitOfTime
from homeassistant.core import HomeAssistant, callback
from homeassistant.helpers.device_registry import DeviceEntryType, DeviceInfo
@ -20,6 +19,7 @@ from homeassistant.helpers.entity_platform import AddEntitiesCallback
from homeassistant.helpers.typing import StateType
from homeassistant.helpers.update_coordinator import CoordinatorEntity
from . import TVTrainConfigEntry
from .const import ATTRIBUTION, DOMAIN
from .coordinator import TrainData, TVDataUpdateCoordinator
@ -106,11 +106,13 @@ SENSOR_TYPES: tuple[TrafikverketSensorEntityDescription, ...] = (
async def async_setup_entry(
hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
hass: HomeAssistant,
entry: TVTrainConfigEntry,
async_add_entities: AddEntitiesCallback,
) -> None:
"""Set up the Trafikverket sensor entry."""
coordinator: TVDataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id]
coordinator = entry.runtime_data
async_add_entities(
[