Move airthings coordinator to separate module (#129158)

homeassistant/components/airthings_ble/__init__.py

@@ -2,75 +2,27 @@
 
 from __future__ import annotations
 
-from datetime import timedelta
-import logging
-
-from airthings_ble import AirthingsBluetoothDeviceData, AirthingsDevice
-from bleak_retry_connector import close_stale_connections_by_address
-
-from homeassistant.components import bluetooth
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import Platform
 from homeassistant.core import HomeAssistant
-from homeassistant.exceptions import ConfigEntryNotReady
-from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed
-from homeassistant.util.unit_system import METRIC_SYSTEM
 
-from .const import DEFAULT_SCAN_INTERVAL, DOMAIN, MAX_RETRIES_AFTER_STARTUP
+from .const import MAX_RETRIES_AFTER_STARTUP
+from .coordinator import AirthingsBLEConfigEntry, AirthingsBLEDataUpdateCoordinator
 
 PLATFORMS: list[Platform] = [Platform.SENSOR]
 
-_LOGGER = logging.getLogger(__name__)
-
-AirthingsBLEDataUpdateCoordinator = DataUpdateCoordinator[AirthingsDevice]
-AirthingsBLEConfigEntry = ConfigEntry[AirthingsBLEDataUpdateCoordinator]
-
 
 async def async_setup_entry(
     hass: HomeAssistant, entry: AirthingsBLEConfigEntry
 ) -> bool:
     """Set up Airthings BLE device from a config entry."""
-    hass.data.setdefault(DOMAIN, {})
+    coordinator = AirthingsBLEDataUpdateCoordinator(hass, entry)
-    address = entry.unique_id
-
-    is_metric = hass.config.units is METRIC_SYSTEM
-    assert address is not None
-
-    await close_stale_connections_by_address(address)
-
-    ble_device = bluetooth.async_ble_device_from_address(hass, address)
-
-    if not ble_device:
-        raise ConfigEntryNotReady(
-            f"Could not find Airthings device with address {address}"
-        )
-
-    airthings = AirthingsBluetoothDeviceData(_LOGGER, is_metric)
-
-    async def _async_update_method() -> AirthingsDevice:
-        """Get data from Airthings BLE."""
-        try:
-            data = await airthings.update_device(ble_device)
-        except Exception as err:
-            raise UpdateFailed(f"Unable to fetch data: {err}") from err
-
-        return data
-
-    coordinator: AirthingsBLEDataUpdateCoordinator = DataUpdateCoordinator(
-        hass,
-        _LOGGER,
-        name=DOMAIN,
-        update_method=_async_update_method,
-        update_interval=timedelta(seconds=DEFAULT_SCAN_INTERVAL),
-    )
-
     await coordinator.async_config_entry_first_refresh()
 
     # Once its setup and we know we are not going to delay
     # the startup of Home Assistant, we can set the max attempts
     # to a higher value. If the first connection attempt fails,
     # Home Assistant's built-in retry logic will take over.
-    airthings.set_max_attempts(MAX_RETRIES_AFTER_STARTUP)
+    coordinator.airthings.set_max_attempts(MAX_RETRIES_AFTER_STARTUP)
 
     entry.runtime_data = coordinator
 

homeassistant/components/airthings_ble/coordinator.py

@@ -0,0 +1,68 @@
+"""The Airthings BLE integration."""
+
+from __future__ import annotations
+
+from datetime import timedelta
+import logging
+
+from airthings_ble import AirthingsBluetoothDeviceData, AirthingsDevice
+from bleak.backends.device import BLEDevice
+from bleak_retry_connector import close_stale_connections_by_address
+
+from homeassistant.components import bluetooth
+from homeassistant.config_entries import ConfigEntry
+from homeassistant.core import HomeAssistant
+from homeassistant.exceptions import ConfigEntryNotReady
+from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed
+from homeassistant.util.unit_system import METRIC_SYSTEM
+
+from .const import DEFAULT_SCAN_INTERVAL, DOMAIN
+
+_LOGGER = logging.getLogger(__name__)
+
+type AirthingsBLEConfigEntry = ConfigEntry[AirthingsBLEDataUpdateCoordinator]
+
+
+class AirthingsBLEDataUpdateCoordinator(DataUpdateCoordinator[AirthingsDevice]):
+    """Class to manage fetching Airthings BLE data."""
+
+    ble_device: BLEDevice
+    config_entry: AirthingsBLEConfigEntry
+
+    def __init__(self, hass: HomeAssistant, entry: AirthingsBLEConfigEntry) -> None:
+        """Initialize the coordinator."""
+        self.airthings = AirthingsBluetoothDeviceData(
+            _LOGGER, hass.config.units is METRIC_SYSTEM
+        )
+        super().__init__(
+            hass,
+            _LOGGER,
+            config_entry=entry,
+            name=DOMAIN,
+            update_interval=timedelta(seconds=DEFAULT_SCAN_INTERVAL),
+        )
+
+    async def _async_setup(self) -> None:
+        """Set up the coordinator."""
+        address = self.config_entry.unique_id
+
+        assert address is not None
+
+        await close_stale_connections_by_address(address)
+
+        ble_device = bluetooth.async_ble_device_from_address(self.hass, address)
+
+        if not ble_device:
+            raise ConfigEntryNotReady(
+                f"Could not find Airthings device with address {address}"
+            )
+        self.ble_device = ble_device
+
+    async def _async_update_data(self) -> AirthingsDevice:
+        """Get data from Airthings BLE."""
+        try:
+            data = await self.airthings.update_device(self.ble_device)
+        except Exception as err:
+            raise UpdateFailed(f"Unable to fetch data: {err}") from err
+
+        return data

homeassistant/components/airthings_ble/sensor.py

@@ -34,8 +34,8 @@ from homeassistant.helpers.typing import StateType
 from homeassistant.helpers.update_coordinator import CoordinatorEntity
 from homeassistant.util.unit_system import METRIC_SYSTEM
 
-from . import AirthingsBLEConfigEntry, AirthingsBLEDataUpdateCoordinator
 from .const import DOMAIN, VOLUME_BECQUEREL, VOLUME_PICOCURIE
+from .coordinator import AirthingsBLEConfigEntry, AirthingsBLEDataUpdateCoordinator
 
 _LOGGER = logging.getLogger(__name__)
 

tests/components/airthings_ble/__init__.py

@@ -49,7 +49,7 @@ def patch_airthings_ble(return_value=AirthingsDevice, side_effect=None):
 def patch_airthings_device_update():
     """Patch airthings-ble device."""
     return patch(
-        "homeassistant.components.airthings_ble.AirthingsBluetoothDeviceData.update_device",
+        "homeassistant.components.airthings_ble.coordinator.AirthingsBluetoothDeviceData.update_device",
         return_value=WAVE_DEVICE_INFO,
     )