Implement opensky data update coordinator (#97925)

Co-authored-by: Robert Resch <robert@resch.dev>

homeassistant/components/opensky/__init__.py

@@ -7,14 +7,17 @@ from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.aiohttp_client import async_get_clientsession
 
-from .const import CLIENT, DOMAIN, PLATFORMS
+from .const import DOMAIN, PLATFORMS
+from .coordinator import OpenSkyDataUpdateCoordinator
 
 
 async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Set up opensky from a config entry."""
 
     client = OpenSky(session=async_get_clientsession(hass))
-    hass.data.setdefault(DOMAIN, {})[entry.entry_id] = {CLIENT: client}
+    coordinator = OpenSkyDataUpdateCoordinator(hass, client)
+    await coordinator.async_config_entry_first_refresh()
+    hass.data.setdefault(DOMAIN, {})[entry.entry_id] = coordinator
 
     await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
 

homeassistant/components/opensky/const.py

@@ -1,11 +1,14 @@
 """OpenSky constants."""
+import logging
+
 from homeassistant.const import Platform
 
+LOGGER = logging.getLogger(__package__)
+
 PLATFORMS = [Platform.SENSOR]
 DEFAULT_NAME = "OpenSky"
 DOMAIN = "opensky"
-CLIENT = "client"
+MANUFACTURER = "OpenSky Network"
-
 CONF_ALTITUDE = "altitude"
 ATTR_ICAO24 = "icao24"
 ATTR_CALLSIGN = "callsign"

homeassistant/components/opensky/coordinator.py

@@ -0,0 +1,116 @@
+"""DataUpdateCoordinator for the OpenSky integration."""
+from __future__ import annotations
+
+from datetime import timedelta
+
+from python_opensky import OpenSky, OpenSkyError, StateVector
+
+from homeassistant.config_entries import ConfigEntry
+from homeassistant.const import (
+    ATTR_LATITUDE,
+    ATTR_LONGITUDE,
+    CONF_LATITUDE,
+    CONF_LONGITUDE,
+    CONF_RADIUS,
+)
+from homeassistant.core import HomeAssistant
+from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed
+
+from .const import (
+    ATTR_ALTITUDE,
+    ATTR_CALLSIGN,
+    ATTR_ICAO24,
+    ATTR_SENSOR,
+    CONF_ALTITUDE,
+    DEFAULT_ALTITUDE,
+    DOMAIN,
+    EVENT_OPENSKY_ENTRY,
+    EVENT_OPENSKY_EXIT,
+    LOGGER,
+)
+
+
+class OpenSkyDataUpdateCoordinator(DataUpdateCoordinator[int]):
+    """An OpenSky Data Update Coordinator."""
+
+    config_entry: ConfigEntry
+
+    def __init__(self, hass: HomeAssistant, opensky: OpenSky) -> None:
+        """Initialize the OpenSky data coordinator."""
+        super().__init__(
+            hass,
+            LOGGER,
+            name=DOMAIN,
+            # OpenSky free user has 400 credits, with 4 credits per API call. 100/24 = ~4 requests per hour
+            update_interval=timedelta(minutes=15),
+        )
+        self._opensky = opensky
+        self._previously_tracked: set[str] | None = None
+        self._bounding_box = OpenSky.get_bounding_box(
+            self.config_entry.data[CONF_LATITUDE],
+            self.config_entry.data[CONF_LONGITUDE],
+            self.config_entry.options[CONF_RADIUS],
+        )
+        self._altitude = self.config_entry.options.get(CONF_ALTITUDE, DEFAULT_ALTITUDE)
+
+    async def _async_update_data(self) -> int:
+        try:
+            response = await self._opensky.get_states(bounding_box=self._bounding_box)
+        except OpenSkyError as exc:
+            raise UpdateFailed from exc
+        currently_tracked = set()
+        flight_metadata: dict[str, StateVector] = {}
+        for flight in response.states:
+            if not flight.callsign:
+                continue
+            callsign = flight.callsign.strip()
+            if callsign:
+                flight_metadata[callsign] = flight
+            else:
+                continue
+            if (
+                flight.longitude is None
+                or flight.latitude is None
+                or flight.on_ground
+                or flight.barometric_altitude is None
+            ):
+                continue
+            altitude = flight.barometric_altitude
+            if altitude > self._altitude and self._altitude != 0:
+                continue
+            currently_tracked.add(callsign)
+        if self._previously_tracked is not None:
+            entries = currently_tracked - self._previously_tracked
+            exits = self._previously_tracked - currently_tracked
+            self._handle_boundary(entries, EVENT_OPENSKY_ENTRY, flight_metadata)
+            self._handle_boundary(exits, EVENT_OPENSKY_EXIT, flight_metadata)
+        self._previously_tracked = currently_tracked
+
+        return len(currently_tracked)
+
+    def _handle_boundary(
+        self, flights: set[str], event: str, metadata: dict[str, StateVector]
+    ) -> None:
+        """Handle flights crossing region boundary."""
+        for flight in flights:
+            if flight in metadata:
+                altitude = metadata[flight].barometric_altitude
+                longitude = metadata[flight].longitude
+                latitude = metadata[flight].latitude
+                icao24 = metadata[flight].icao24
+            else:
+                # Assume Flight has landed if missing.
+                altitude = 0
+                longitude = None
+                latitude = None
+                icao24 = None
+
+            data = {
+                ATTR_CALLSIGN: flight,
+                ATTR_ALTITUDE: altitude,
+                ATTR_SENSOR: self.config_entry.title,
+                ATTR_LONGITUDE: longitude,
+                ATTR_LATITUDE: latitude,
+                ATTR_ICAO24: icao24,
+            }
+            self.hass.bus.fire(event, data)

homeassistant/components/opensky/sensor.py

@@ -1,16 +1,15 @@
 """Sensor for the Open Sky Network."""
 from __future__ import annotations
 
-from datetime import timedelta
-
-from python_opensky import BoundingBox, OpenSky, StateVector
 import voluptuous as vol
 
-from homeassistant.components.sensor import PLATFORM_SCHEMA, SensorEntity
+from homeassistant.components.sensor import (
+    PLATFORM_SCHEMA,
+    SensorEntity,
+    SensorStateClass,
+)
 from homeassistant.config_entries import SOURCE_IMPORT, ConfigEntry
 from homeassistant.const import (
-    ATTR_LATITUDE,
-    ATTR_LONGITUDE,
     CONF_LATITUDE,
     CONF_LONGITUDE,
     CONF_NAME,
@@ -18,26 +17,20 @@ from homeassistant.const import (
 )
 from homeassistant.core import DOMAIN as HOMEASSISTANT_DOMAIN, HomeAssistant
 import homeassistant.helpers.config_validation as cv
+from homeassistant.helpers.device_registry import DeviceEntryType
+from homeassistant.helpers.entity import DeviceInfo
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.issue_registry import IssueSeverity, async_create_issue
 from homeassistant.helpers.typing import ConfigType, DiscoveryInfoType
+from homeassistant.helpers.update_coordinator import CoordinatorEntity
 
 from .const import (
-    ATTR_ALTITUDE,
-    ATTR_CALLSIGN,
-    ATTR_ICAO24,
-    ATTR_SENSOR,
-    CLIENT,
     CONF_ALTITUDE,
     DEFAULT_ALTITUDE,
     DOMAIN,
-    EVENT_OPENSKY_ENTRY,
+    MANUFACTURER,
-    EVENT_OPENSKY_EXIT,
 )
-
+from .coordinator import OpenSkyDataUpdateCoordinator
-# OpenSky free user has 400 credits, with 4 credits per API call. 100/24 = ~4 requests per hour
-SCAN_INTERVAL = timedelta(minutes=15)
-
 
 PLATFORM_SCHEMA = PLATFORM_SCHEMA.extend(
     {
@@ -87,125 +80,45 @@ async def async_setup_entry(
 ) -> None:
     """Initialize the entries."""
 
-    opensky = hass.data[DOMAIN][entry.entry_id][CLIENT]
+    coordinator = hass.data[DOMAIN][entry.entry_id]
-    bounding_box = OpenSky.get_bounding_box(
-        entry.data[CONF_LATITUDE],
-        entry.data[CONF_LONGITUDE],
-        entry.options[CONF_RADIUS],
-    )
     async_add_entities(
         [
             OpenSkySensor(
-                entry.title,
+                coordinator,
-                opensky,
+                entry,
-                bounding_box,
-                entry.options.get(CONF_ALTITUDE, DEFAULT_ALTITUDE),
-                entry.entry_id,
             )
         ],
-        True,
     )
 
 
-class OpenSkySensor(SensorEntity):
+class OpenSkySensor(CoordinatorEntity[OpenSkyDataUpdateCoordinator], SensorEntity):
     """Open Sky Network Sensor."""
 
     _attr_attribution = (
         "Information provided by the OpenSky Network (https://opensky-network.org)"
     )
+    _attr_has_entity_name = True
+    _attr_name = None
+    _attr_icon = "mdi:airplane"
+    _attr_native_unit_of_measurement = "flights"
+    _attr_state_class = SensorStateClass.MEASUREMENT
 
     def __init__(
         self,
-        name: str,
+        coordinator: OpenSkyDataUpdateCoordinator,
-        opensky: OpenSky,
+        config_entry: ConfigEntry,
-        bounding_box: BoundingBox,
-        altitude: float,
-        entry_id: str,
     ) -> None:
         """Initialize the sensor."""
-        self._altitude = altitude
+        super().__init__(coordinator)
-        self._state = 0
+        self._attr_unique_id = f"{config_entry.entry_id}_opensky"
-        self._name = name
+        self._attr_device_info = DeviceInfo(
-        self._previously_tracked: set[str] = set()
+            identifiers={(DOMAIN, f"{coordinator.config_entry.entry_id}")},
-        self._opensky = opensky
+            manufacturer=MANUFACTURER,
-        self._bounding_box = bounding_box
+            name=config_entry.title,
-        self._attr_unique_id = f"{entry_id}_opensky"
+            entry_type=DeviceEntryType.SERVICE,
-
+        )
-    @property
-    def name(self) -> str:
-        """Return the name of the sensor."""
-        return self._name
 
     @property
     def native_value(self) -> int:
         """Return the state of the sensor."""
-        return self._state
+        return self.coordinator.data
-
-    def _handle_boundary(
-        self, flights: set[str], event: str, metadata: dict[str, StateVector]
-    ) -> None:
-        """Handle flights crossing region boundary."""
-        for flight in flights:
-            if flight in metadata:
-                altitude = metadata[flight].barometric_altitude
-                longitude = metadata[flight].longitude
-                latitude = metadata[flight].latitude
-                icao24 = metadata[flight].icao24
-            else:
-                # Assume Flight has landed if missing.
-                altitude = 0
-                longitude = None
-                latitude = None
-                icao24 = None
-
-            data = {
-                ATTR_CALLSIGN: flight,
-                ATTR_ALTITUDE: altitude,
-                ATTR_SENSOR: self._name,
-                ATTR_LONGITUDE: longitude,
-                ATTR_LATITUDE: latitude,
-                ATTR_ICAO24: icao24,
-            }
-            self.hass.bus.fire(event, data)
-
-    async def async_update(self) -> None:
-        """Update device state."""
-        currently_tracked = set()
-        flight_metadata: dict[str, StateVector] = {}
-        response = await self._opensky.get_states(bounding_box=self._bounding_box)
-        for flight in response.states:
-            if not flight.callsign:
-                continue
-            callsign = flight.callsign.strip()
-            if callsign != "":
-                flight_metadata[callsign] = flight
-            else:
-                continue
-            if (
-                flight.longitude is None
-                or flight.latitude is None
-                or flight.on_ground
-                or flight.barometric_altitude is None
-            ):
-                continue
-            altitude = flight.barometric_altitude
-            if altitude > self._altitude and self._altitude != 0:
-                continue
-            currently_tracked.add(callsign)
-        if self._previously_tracked is not None:
-            entries = currently_tracked - self._previously_tracked
-            exits = self._previously_tracked - currently_tracked
-            self._handle_boundary(entries, EVENT_OPENSKY_ENTRY, flight_metadata)
-            self._handle_boundary(exits, EVENT_OPENSKY_EXIT, flight_metadata)
-        self._state = len(currently_tracked)
-        self._previously_tracked = currently_tracked
-
-    @property
-    def native_unit_of_measurement(self) -> str:
-        """Return the unit of measurement."""
-        return "flights"
-
-    @property
-    def icon(self) -> str:
-        """Return the icon."""
-        return "mdi:airplane"

tests/components/opensky/snapshots/test_sensor.ambr

@@ -5,6 +5,7 @@
       'attribution': 'Information provided by the OpenSky Network (https://opensky-network.org)',
       'friendly_name': 'OpenSky',
       'icon': 'mdi:airplane',
+      'state_class': <SensorStateClass.MEASUREMENT: 'measurement'>,
       'unit_of_measurement': 'flights',
     }),
     'context': <ANY>,
@@ -20,6 +21,7 @@
       'attribution': 'Information provided by the OpenSky Network (https://opensky-network.org)',
       'friendly_name': 'OpenSky',
       'icon': 'mdi:airplane',
+      'state_class': <SensorStateClass.MEASUREMENT: 'measurement'>,
       'unit_of_measurement': 'flights',
     }),
     'context': <ANY>,

tests/components/opensky/test_init.py

@@ -1,8 +1,14 @@
 """Test OpenSky component setup process."""
 from __future__ import annotations
 
+from unittest.mock import patch
+
+from python_opensky import OpenSkyError
+
 from homeassistant.components.opensky.const import DOMAIN
+from homeassistant.config_entries import ConfigEntryState
 from homeassistant.core import HomeAssistant
+from homeassistant.setup import async_setup_component
 
 from .conftest import ComponentSetup
 
@@ -26,3 +32,19 @@ async def test_load_unload_entry(
 
     state = hass.states.get("sensor.opensky")
     assert not state
+
+
+async def test_load_entry_failure(
+    hass: HomeAssistant,
+    config_entry: MockConfigEntry,
+) -> None:
+    """Test failure while loading."""
+    config_entry.add_to_hass(hass)
+    with patch(
+        "python_opensky.OpenSky.get_states",
+        side_effect=OpenSkyError(),
+    ):
+        assert await async_setup_component(hass, DOMAIN, {})
+        await hass.async_block_till_done()
+    entry = hass.config_entries.async_entries(DOMAIN)[0]
+    assert entry.state == ConfigEntryState.SETUP_RETRY