From 98c81fa2af644919fca7ff4f15da994624ba0ff5 Mon Sep 17 00:00:00 2001 From: epenet <6771947+epenet@users.noreply.github.com> Date: Sat, 26 Oct 2024 02:29:57 +0200 Subject: [PATCH] Move airthings coordinator to separate module (#129158) --- .../components/airthings_ble/__init__.py | 56 ++------------- .../components/airthings_ble/coordinator.py | 68 +++++++++++++++++++ .../components/airthings_ble/sensor.py | 2 +- tests/components/airthings_ble/__init__.py | 2 +- 4 files changed, 74 insertions(+), 54 deletions(-) create mode 100644 homeassistant/components/airthings_ble/coordinator.py diff --git a/homeassistant/components/airthings_ble/__init__.py b/homeassistant/components/airthings_ble/__init__.py index 79384eed4ef..1c3c6084739 100644 --- a/homeassistant/components/airthings_ble/__init__.py +++ b/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, {}) - 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), - ) - + coordinator = AirthingsBLEDataUpdateCoordinator(hass, entry) 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 diff --git a/homeassistant/components/airthings_ble/coordinator.py b/homeassistant/components/airthings_ble/coordinator.py new file mode 100644 index 00000000000..81009dcea81 --- /dev/null +++ b/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 diff --git a/homeassistant/components/airthings_ble/sensor.py b/homeassistant/components/airthings_ble/sensor.py index b1ae7d533d8..0dfd82a38c4 100644 --- a/homeassistant/components/airthings_ble/sensor.py +++ b/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__) diff --git a/tests/components/airthings_ble/__init__.py b/tests/components/airthings_ble/__init__.py index a736fa979e9..add21b1067f 100644 --- a/tests/components/airthings_ble/__init__.py +++ b/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, )