feat: fix some issues and move to bluetooth update coordinator

hunterjm · hunterjm · commit 9c75c6fee140 · 2023-07-15T16:10:22.000-04:00
diff --git a/custom_components/ac_infinity/__init__.py b/custom_components/ac_infinity/__init__.py
@@ -1,27 +1,22 @@
 """The ac_infinity integration."""
 from __future__ import annotations
 
-import asyncio
-from datetime import timedelta
 import logging
 
-import async_timeout
 from ac_infinity_ble import ACInfinityController, DeviceInfo
 
 from homeassistant.components import bluetooth
-from homeassistant.components.bluetooth.match import ADDRESS, BluetoothCallbackMatcher
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import (
     CONF_ADDRESS,
     CONF_SERVICE_DATA,
-    EVENT_HOMEASSISTANT_STOP,
     Platform,
 )
-from homeassistant.core import Event, HomeAssistant, callback
+from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import ConfigEntryNotReady
-from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed
 
-from .const import BLEAK_EXCEPTIONS, DEVICE_TIMEOUT, DOMAIN, UPDATE_SECONDS
+from .const import DOMAIN
+from .coordinator import ACInfinityDataUpdateCoordinator
 from .models import ACInfinityData
 
 PLATFORMS: list[Platform] = [Platform.SENSOR, Platform.FAN]
@@ -42,78 +37,19 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     if type(device_info) is dict:
         device_info = DeviceInfo(**entry.data[CONF_SERVICE_DATA])
     controller = ACInfinityController(ble_device, device_info)
-
-    @callback
-    def _async_update_ble(
-        service_info: bluetooth.BluetoothServiceInfoBleak,
-        change: bluetooth.BluetoothChange,
-    ) -> None:
-        """Update from a ble callback."""
-        controller.set_ble_device_and_advertisement_data(
-            service_info.device, service_info.advertisement
-        )
-
-    entry.async_on_unload(
-        bluetooth.async_register_callback(
-            hass,
-            _async_update_ble,
-            BluetoothCallbackMatcher({ADDRESS: address}),
-            bluetooth.BluetoothScanningMode.PASSIVE,
-        )
-    )
-
-    async def _async_update():
-        """Update the device state."""
-        try:
-            await controller.update()
-            await controller.stop()
-        except BLEAK_EXCEPTIONS as ex:
-            raise UpdateFailed(str(ex)) from ex
-
-    startup_event = asyncio.Event()
-    cancel_first_update = controller.register_callback(lambda *_: startup_event.set())
-    coordinator = DataUpdateCoordinator(
-        hass,
-        _LOGGER,
-        name=controller.name,
-        update_method=_async_update,
-        update_interval=timedelta(seconds=UPDATE_SECONDS),
-    )
-
-    try:
-        await coordinator.async_config_entry_first_refresh()
-    except ConfigEntryNotReady:
-        cancel_first_update()
-        raise
-
-    try:
-        async with async_timeout.timeout(DEVICE_TIMEOUT):
-            await startup_event.wait()
-    except asyncio.TimeoutError as ex:
-        raise ConfigEntryNotReady(
-            "Unable to communicate with the device; "
-            f"Try moving the Bluetooth adapter closer to {controller.name}"
-        ) from ex
-    finally:
-        cancel_first_update()
+    coordinator = ACInfinityDataUpdateCoordinator(hass, _LOGGER, ble_device, controller)
 
     hass.data.setdefault(DOMAIN, {})[entry.entry_id] = ACInfinityData(
         entry.title, controller, coordinator
     )
 
-    await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
-    entry.async_on_unload(entry.add_update_listener(_async_update_listener))
+    entry.async_on_unload(coordinator.async_start())
+    if not await coordinator.async_wait_ready():
+        raise ConfigEntryNotReady(f"{address} is not advertising state")
 
-    async def _async_stop(event: Event) -> None:
-        """Close the connection."""
-        try:
-            await controller.stop()
-        except BLEAK_EXCEPTIONS:
-            pass
+    entry.async_on_unload(entry.add_update_listener(_async_update_listener))
+    await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
 
-    entry.async_on_unload(
-        hass.bus.async_listen_once(EVENT_HOMEASSISTANT_STOP, _async_stop)
-    )
     return True
 
 
@@ -127,10 +63,6 @@ async def _async_update_listener(hass: HomeAssistant, entry: ConfigEntry) -> Non
 async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Unload a config entry."""
     if unload_ok := await hass.config_entries.async_unload_platforms(entry, PLATFORMS):
-        data: ACInfinityData = hass.data[DOMAIN].pop(entry.entry_id)
-        try:
-            await data.device.stop()
-        except BLEAK_EXCEPTIONS:
-            pass
+        hass.data[DOMAIN].pop(entry.entry_id)
 
     return unload_ok
diff --git a/custom_components/ac_infinity/const.py b/custom_components/ac_infinity/const.py
@@ -6,6 +6,6 @@
 DEVICE_TIMEOUT = 30
 UPDATE_SECONDS = 15
 
-BLEAK_EXCEPTIONS = (AttributeError, BleakError)
+BLEAK_EXCEPTIONS = (AttributeError, BleakError, TimeoutError)
 
 DEVICE_MODEL = {1: "Controller 67", 7: "Controller 69", 11: "Controller 69 Pro"}
diff --git a/custom_components/ac_infinity/coordinator.py b/custom_components/ac_infinity/coordinator.py
@@ -0,0 +1,104 @@
+"""AC Infinity Coordinator."""
+from __future__ import annotations
+
+import asyncio
+import contextlib
+import logging
+
+from ac_infinity_ble import ACInfinityController
+import async_timeout
+from bleak.backends.device import BLEDevice
+
+from homeassistant.components import bluetooth
+from homeassistant.components.bluetooth.active_update_coordinator import (
+    ActiveBluetoothDataUpdateCoordinator,
+)
+from homeassistant.core import CoreState, HomeAssistant, callback
+
+
+DEVICE_STARTUP_TIMEOUT = 30
+
+
+class ACInfinityDataUpdateCoordinator(ActiveBluetoothDataUpdateCoordinator[None]):
+    """Class to manage fetching switchbot data."""
+
+    def __init__(
+        self,
+        hass: HomeAssistant,
+        logger: logging.Logger,
+        ble_device: BLEDevice,
+        controller: ACInfinityController,
+    ) -> None:
+        """Initialize global switchbot data updater."""
+        super().__init__(
+            hass=hass,
+            logger=logger,
+            address=ble_device.address,
+            needs_poll_method=self._needs_poll,
+            poll_method=self._async_update,
+            mode=bluetooth.BluetoothScanningMode.ACTIVE,
+            connectable=True,
+        )
+        self.ble_device = ble_device
+        self.controller = controller
+        self._ready_event = asyncio.Event()
+        self._was_unavailable = True
+
+    @callback
+    def _needs_poll(
+        self,
+        service_info: bluetooth.BluetoothServiceInfoBleak,
+        seconds_since_last_poll: float | None,
+    ) -> bool:
+        # Only poll if hass is running, we need to poll,
+        # and we actually have a way to connect to the device
+        return (
+            self.hass.state == CoreState.running
+            and (seconds_since_last_poll is None or seconds_since_last_poll > 30)
+            and bool(
+                bluetooth.async_ble_device_from_address(
+                    self.hass, service_info.device.address, connectable=True
+                )
+            )
+        )
+
+    async def _async_update(
+        self, service_info: bluetooth.BluetoothServiceInfoBleak
+    ) -> None:
+        """Poll the device."""
+        await self.controller.update()
+
+    @callback
+    def _async_handle_unavailable(
+        self, service_info: bluetooth.BluetoothServiceInfoBleak
+    ) -> None:
+        """Handle the device going unavailable."""
+        super()._async_handle_unavailable(service_info)
+        self._was_unavailable = True
+
+    @callback
+    def _async_handle_bluetooth_event(
+        self,
+        service_info: bluetooth.BluetoothServiceInfoBleak,
+        change: bluetooth.BluetoothChange,
+    ) -> None:
+        """Handle a Bluetooth event."""
+        self.ble_device = service_info.device
+        self.controller.set_ble_device_and_advertisement_data(
+            service_info.device, service_info.advertisement
+        )
+        if self.controller.name:
+            self._ready_event.set()
+        self.logger.debug(
+            "%s: AC Infinity data: %s", self.ble_device.address, self.controller.state
+        )
+        self._was_unavailable = False
+        super()._async_handle_bluetooth_event(service_info, change)
+
+    async def async_wait_ready(self) -> bool:
+        """Wait for the device to be ready."""
+        with contextlib.suppress(asyncio.TimeoutError):
+            async with async_timeout.timeout(DEVICE_STARTUP_TIMEOUT):
+                await self._ready_event.wait()
+                return True
+        return False
diff --git a/custom_components/ac_infinity/fan.py b/custom_components/ac_infinity/fan.py
@@ -8,22 +8,22 @@
 
 from homeassistant.components.fan import FanEntity, FanEntityFeature
 
+from homeassistant.components.bluetooth.passive_update_coordinator import (
+    PassiveBluetoothCoordinatorEntity,
+)
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers import device_registry as dr
 from homeassistant.helpers.entity import DeviceInfo
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
-from homeassistant.helpers.update_coordinator import (
-    CoordinatorEntity,
-    DataUpdateCoordinator,
-)
 from homeassistant.util.percentage import (
     int_states_in_range,
     ranged_value_to_percentage,
     percentage_to_ranged_value,
 )
 
 from .const import DEVICE_MODEL, DOMAIN
+from .coordinator import ACInfinityDataUpdateCoordinator
 from .models import ACInfinityData
 
 SPEED_RANGE = (1, 10)
@@ -39,15 +39,17 @@ async def async_setup_entry(
     async_add_entities([ACInfinityFan(data.coordinator, data.device, entry.title)])
 
 
-class ACInfinityFan(CoordinatorEntity, FanEntity):
+class ACInfinityFan(
+    PassiveBluetoothCoordinatorEntity[ACInfinityDataUpdateCoordinator], FanEntity
+):
     """Representation of AC Infinity sensor."""
 
     _attr_speed_count = int_states_in_range(SPEED_RANGE)
     _attr_supported_features = FanEntityFeature.SET_SPEED
 
     def __init__(
         self,
-        coordinator: DataUpdateCoordinator,
+        coordinator: ACInfinityDataUpdateCoordinator,
         device: ACInfinityController,
         name: str,
     ) -> None:
@@ -72,7 +74,6 @@ async def async_set_percentage(self, percentage: int) -> None:
             speed = math.ceil(percentage_to_ranged_value(SPEED_RANGE, percentage))
 
         await self._device.set_speed(speed)
-        self._async_update_attrs()
 
     async def async_turn_on(
         self,
@@ -85,12 +86,10 @@ async def async_turn_on(
         if percentage is not None:
             speed = math.ceil(percentage_to_ranged_value(SPEED_RANGE, percentage))
         await self._device.turn_on(speed)
-        self._async_update_attrs()
 
     async def async_turn_off(self, **kwargs: Any) -> None:
         """Turn off the fan."""
         await self._device.turn_off()
-        self._async_update_attrs()
 
     @callback
     def _async_update_attrs(self) -> None:
diff --git a/custom_components/ac_infinity/manifest.json b/custom_components/ac_infinity/manifest.json
@@ -19,6 +19,6 @@
   "issue_tracker": "https://github.com/hunterjm/ac-infinity-hacs/issues",
   "iot_class": "local_push",
   "requirements": [
-    "ac-infinity-ble==0.4.1"
+    "ac-infinity-ble==0.4.2"
   ]
 }
diff --git a/custom_components/ac_infinity/models.py b/custom_components/ac_infinity/models.py
@@ -5,7 +5,7 @@
 
 from ac_infinity_ble import ACInfinityController
 
-from homeassistant.helpers.update_coordinator import DataUpdateCoordinator
+from .coordinator import ACInfinityDataUpdateCoordinator
 
 
 @dataclass
@@ -14,4 +14,4 @@ class ACInfinityData:
 
     title: str
     device: ACInfinityController
-    coordinator: DataUpdateCoordinator
+    coordinator: ACInfinityDataUpdateCoordinator
diff --git a/custom_components/ac_infinity/sensor.py b/custom_components/ac_infinity/sensor.py
@@ -10,19 +10,18 @@
     SensorStateClass,
 )
 
+from homeassistant.components.bluetooth.passive_update_coordinator import (
+    PassiveBluetoothCoordinatorEntity,
+)
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import PERCENTAGE, UnitOfPressure, UnitOfTemperature
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers import device_registry as dr
 from homeassistant.helpers.entity import DeviceInfo
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
-from homeassistant.helpers.typing import UndefinedType
-from homeassistant.helpers.update_coordinator import (
-    CoordinatorEntity,
-    DataUpdateCoordinator,
-)
 
 from .const import DEVICE_MODEL, DOMAIN
+from .coordinator import ACInfinityDataUpdateCoordinator
 from .models import ACInfinityData
 
 
@@ -33,21 +32,23 @@ async def async_setup_entry(
 ) -> None:
     """Set up the light platform for LEDBLE."""
     data: ACInfinityData = hass.data[DOMAIN][entry.entry_id]
-    async_add_entities(
-        [
-            TemperatureSensor(data.coordinator, data.device, entry.title),
-            HumiditySensor(data.coordinator, data.device, entry.title),
-            VpdSensor(data.coordinator, data.device, entry.title),
-        ]
-    )
-
-
-class ACInfinitySensor(CoordinatorEntity, SensorEntity):
+    entities = [
+        TemperatureSensor(data.coordinator, data.device, entry.title),
+        HumiditySensor(data.coordinator, data.device, entry.title),
+    ]
+    if data.device.state.version >= 3 and data.device.state.type in [7, 9, 11, 12]:
+        entities.append(VpdSensor(data.coordinator, data.device, entry.title))
+    async_add_entities(entities)
+
+
+class ACInfinitySensor(
+    PassiveBluetoothCoordinatorEntity[ACInfinityDataUpdateCoordinator], SensorEntity
+):
     """Representation of AC Infinity sensor."""
 
     def __init__(
         self,
-        coordinator: DataUpdateCoordinator,
+        coordinator: ACInfinityDataUpdateCoordinator,
         device: ACInfinityController,
         name: str,
     ) -> None:

Original file line number	Diff line number	Diff line change
`@@ -19,6 +19,6 @@`
`19`	`19`	`"issue_tracker": "https://github.com/hunterjm/ac-infinity-hacs/issues",`
`20`	`20`	`"iot_class": "local_push",`
`21`	`21`	`"requirements": [`
`22`		`- "ac-infinity-ble==0.4.1"`
	`22`	`+ "ac-infinity-ble==0.4.2"`
`23`	`23`	`]`
`24`	`24`	`}`