diff --git a/homeassistant/components/london_underground/sensor.py b/homeassistant/components/london_underground/sensor.py index a4cc66a8447..96ff9bc5056 100644 --- a/homeassistant/components/london_underground/sensor.py +++ b/homeassistant/components/london_underground/sensor.py @@ -2,17 +2,28 @@ from __future__ import annotations from datetime import timedelta +import logging +import async_timeout from london_tube_status import TubeData import voluptuous as vol from homeassistant.components.sensor import PLATFORM_SCHEMA, SensorEntity from homeassistant.const import ATTR_ATTRIBUTION from homeassistant.core import HomeAssistant +from homeassistant.exceptions import PlatformNotReady from homeassistant.helpers.aiohttp_client import async_get_clientsession import homeassistant.helpers.config_validation as cv from homeassistant.helpers.entity_platform import AddEntitiesCallback from homeassistant.helpers.typing import ConfigType, DiscoveryInfoType +from homeassistant.helpers.update_coordinator import ( + CoordinatorEntity, + DataUpdateCoordinator, +) + +_LOGGER = logging.getLogger(__name__) + +DOMAIN = "london_underground" ATTRIBUTION = "Powered by TfL Open Data" @@ -55,24 +66,46 @@ async def async_setup_platform( session = async_get_clientsession(hass) data = TubeData(session) - await data.update() + coordinator = LondonTubeCoordinator(hass, data) + + await coordinator.async_refresh() + + if not coordinator.last_update_success: + raise PlatformNotReady sensors = [] for line in config[CONF_LINE]: - sensors.append(LondonTubeSensor(line, data)) + sensors.append(LondonTubeSensor(coordinator, line)) - async_add_entities(sensors, True) + async_add_entities(sensors) -class LondonTubeSensor(SensorEntity): +class LondonTubeCoordinator(DataUpdateCoordinator): + """London Underground sensor coordinator.""" + + def __init__(self, hass, data): + """Initialize coordinator.""" + super().__init__( + hass, + _LOGGER, + name=DOMAIN, + update_interval=SCAN_INTERVAL, + ) + self._data = data + + async def _async_update_data(self): + async with async_timeout.timeout(10): + await self._data.update() + return self._data.data + + +class LondonTubeSensor(CoordinatorEntity[LondonTubeCoordinator], SensorEntity): """Sensor that reads the status of a line from Tube Data.""" - def __init__(self, name, data): + def __init__(self, coordinator, name): """Initialize the London Underground sensor.""" - self._data = data - self._description = None + super().__init__(coordinator) self._name = name - self._state = None self.attrs = {ATTR_ATTRIBUTION: ATTRIBUTION} @property @@ -83,7 +116,7 @@ class LondonTubeSensor(SensorEntity): @property def native_value(self): """Return the state of the sensor.""" - return self._state + return self.coordinator.data[self.name]["State"] @property def icon(self): @@ -93,11 +126,5 @@ class LondonTubeSensor(SensorEntity): @property def extra_state_attributes(self): """Return other details about the sensor state.""" - self.attrs["Description"] = self._description + self.attrs["Description"] = self.coordinator.data[self.name]["Description"] return self.attrs - - async def async_update(self): - """Update the sensor.""" - await self._data.update() - self._state = self._data.data[self.name]["State"] - self._description = self._data.data[self.name]["Description"]