Add better connection management for Idasen Desk (home-assistant#102135)

homeassistant/components/idasen_desk/__init__.py

@@ -4,8 +4,9 @@
 import logging
 
 from attr import dataclass
-from bleak import BleakError
+from bleak.exc import BleakError
 from idasen_ha import Desk
+from idasen_ha.errors import AuthFailedError
 
 from homeassistant.components import bluetooth
 from homeassistant.config_entries import ConfigEntry
@@ -15,7 +16,7 @@
     EVENT_HOMEASSISTANT_STOP,
     Platform,
 )
-from homeassistant.core import Event, HomeAssistant
+from homeassistant.core import Event, HomeAssistant, callback
 from homeassistant.exceptions import ConfigEntryNotReady
 from homeassistant.helpers import device_registry as dr
 from homeassistant.helpers.device_registry import DeviceInfo
@@ -28,49 +29,92 @@
 _LOGGER = logging.getLogger(__name__)
 
 
+class IdasenDeskCoordinator(DataUpdateCoordinator):
+    """Class to manage updates for the Idasen Desk."""
+
+    def __init__(
+        self,
+        hass: HomeAssistant,
+        logger: logging.Logger,
+        name: str,
+        address: str,
+    ) -> None:
+        """Init IdasenDeskCoordinator."""
+
+        super().__init__(hass, logger, name=name)
+        self._address = address
+        self._expected_connected = False
+
+        self.desk = Desk(self.async_set_updated_data)
+
+    async def async_connect(self) -> bool:
+        """Connect to desk."""
+        _LOGGER.debug("Trying to connect %s", self._address)
+        ble_device = bluetooth.async_ble_device_from_address(
+            self.hass, self._address, connectable=True
+        )
+        if ble_device is None:
+            return False
+        self._expected_connected = True
+        await self.desk.connect(ble_device)
+        return True
+
+    async def async_disconnect(self) -> None:
+        """Disconnect from desk."""
+        _LOGGER.debug("Disconnecting from %s", self._address)
+        self._expected_connected = False
+        await self.desk.disconnect()
+
+    @callback
+    def async_set_updated_data(self, data: int | None) -> None:
+        """Handle data update."""
+        if self._expected_connected:
+            if not self.desk.is_connected:
+                _LOGGER.debug("Desk disconnected. Reconnecting")
+                self.hass.async_create_task(self.async_connect())
+        elif self.desk.is_connected:
+            _LOGGER.warning("Desk is connected but should not be. Disconnecting")
+            self.hass.async_create_task(self.desk.disconnect())
+        return super().async_set_updated_data(data)
+
+
 @dataclass
 class DeskData:
     """Data for the Idasen Desk integration."""
 
-    desk: Desk
     address: str
     device_info: DeviceInfo
-    coordinator: DataUpdateCoordinator
+    coordinator: IdasenDeskCoordinator
 
 
 async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Set up IKEA Idasen from a config entry."""
     address: str = entry.data[CONF_ADDRESS].upper()
 
-    coordinator: DataUpdateCoordinator = DataUpdateCoordinator(
-        hass,
-        _LOGGER,
-        name=entry.title,
+    coordinator: IdasenDeskCoordinator = IdasenDeskCoordinator(
+        hass, _LOGGER, entry.title, address
     )
 
-    desk = Desk(coordinator.async_set_updated_data)
     device_info = DeviceInfo(
         name=entry.title,
         connections={(dr.CONNECTION_BLUETOOTH, address)},
     )
     hass.data.setdefault(DOMAIN, {})[entry.entry_id] = DeskData(
-        desk, address, device_info, coordinator
+        address, device_info, coordinator
     )
 
-    ble_device = bluetooth.async_ble_device_from_address(
-        hass, address, connectable=True
-    )
     try:
-        await desk.connect(ble_device)
-    except (TimeoutError, BleakError) as ex:
+        if not await coordinator.async_connect():
+            raise ConfigEntryNotReady(f"Unable to connect to desk {address}")
+    except (AuthFailedError, TimeoutError, BleakError, Exception) as ex:
         raise ConfigEntryNotReady(f"Unable to connect to desk {address}") from ex
 
     await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
     entry.async_on_unload(entry.add_update_listener(_async_update_listener))
 
     async def _async_stop(event: Event) -> None:
         """Close the connection."""
-        await desk.disconnect()
+        await coordinator.async_disconnect()
 
     entry.async_on_unload(
         hass.bus.async_listen_once(EVENT_HOMEASSISTANT_STOP, _async_stop)
@@ -89,7 +133,7 @@ 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: DeskData = hass.data[DOMAIN].pop(entry.entry_id)
-        await data.desk.disconnect()
+        await data.coordinator.async_disconnect()
         bluetooth.async_rediscover_address(hass, data.address)
 
     return unload_ok

homeassistant/components/idasen_desk/config_flow.py

@@ -6,7 +6,8 @@
 
 from bleak.exc import BleakError
 from bluetooth_data_tools import human_readable_name
-from idasen_ha import AuthFailedError, Desk
+from idasen_ha import Desk
+from idasen_ha.errors import AuthFailedError
 import voluptuous as vol
 
 from homeassistant import config_entries
@@ -61,9 +62,9 @@ async def async_step_user(
             )
             self._abort_if_unique_id_configured()
 
-            desk = Desk(None)
+            desk = Desk(None, monitor_height=False)
             try:
-                await desk.connect(discovery_info.device, monitor_height=False)
+                await desk.connect(discovery_info.device, auto_reconnect=False)
             except AuthFailedError as err:
                 _LOGGER.exception("AuthFailedError", exc_info=err)
                 errors["base"] = "auth_failed"

homeassistant/components/idasen_desk/cover.py

@@ -1,11 +1,8 @@
 """Idasen Desk integration cover platform."""
 from __future__ import annotations
 
-import logging
 from typing import Any
 
-from idasen_ha import Desk
-
 from homeassistant.components.cover import (
     ATTR_POSITION,
     CoverDeviceClass,
@@ -17,16 +14,11 @@
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.device_registry import DeviceInfo
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
-from homeassistant.helpers.update_coordinator import (
-    CoordinatorEntity,
-    DataUpdateCoordinator,
-)
+from homeassistant.helpers.update_coordinator import CoordinatorEntity
 
-from . import DeskData
+from . import DeskData, IdasenDeskCoordinator
 from .const import DOMAIN
 
-_LOGGER = logging.getLogger(__name__)
-
 
 async def async_setup_entry(
     hass: HomeAssistant,
@@ -36,7 +28,7 @@ async def async_setup_entry(
     """Set up the cover platform for Idasen Desk."""
     data: DeskData = hass.data[DOMAIN][entry.entry_id]
     async_add_entities(
-        [IdasenDeskCover(data.desk, data.address, data.device_info, data.coordinator)]
+        [IdasenDeskCover(data.address, data.device_info, data.coordinator)]
     )
 
 
@@ -54,14 +46,13 @@ class IdasenDeskCover(CoordinatorEntity, CoverEntity):
 
     def __init__(
         self,
-        desk: Desk,
         address: str,
         device_info: DeviceInfo,
-        coordinator: DataUpdateCoordinator,
+        coordinator: IdasenDeskCoordinator,
     ) -> None:
         """Initialize an Idasen Desk cover."""
         super().__init__(coordinator)
-        self._desk = desk
+        self._desk = coordinator.desk
         self._attr_name = device_info[ATTR_NAME]
         self._attr_unique_id = address
         self._attr_device_info = device_info

homeassistant/components/idasen_desk/manifest.json

@@ -11,5 +11,5 @@
   "dependencies": ["bluetooth_adapters"],
   "documentation": "https://www.home-assistant.io/integrations/idasen_desk",
   "iot_class": "local_push",
-  "requirements": ["idasen-ha==1.4.1"]
+  "requirements": ["idasen-ha==2.3"]
 }

requirements_all.txt

@@ -1051,7 +1051,7 @@ ical==5.0.1
 icmplib==3.0
 
 # homeassistant.components.idasen_desk
-idasen-ha==1.4.1
+idasen-ha==2.3
 
 # homeassistant.components.network
 ifaddr==0.2.0

requirements_test_all.txt

@@ -831,7 +831,7 @@ ical==5.0.1
 icmplib==3.0
 
 # homeassistant.components.idasen_desk
-idasen-ha==1.4.1
+idasen-ha==2.3
 
 # homeassistant.components.network
 ifaddr==0.2.0

tests/components/idasen_desk/conftest.py

@@ -10,6 +10,10 @@
 @pytest.fixture(autouse=True)
 def mock_bluetooth(enable_bluetooth):
     """Auto mock bluetooth."""
+    with mock.patch(
+        "homeassistant.components.idasen_desk.bluetooth.async_ble_device_from_address"
+    ):
+        yield MagicMock()
 
 
 @pytest.fixture(autouse=False)
@@ -18,14 +22,22 @@ def mock_desk_api():
     with mock.patch("homeassistant.components.idasen_desk.Desk") as desk_patched:
         mock_desk = MagicMock()
 
-        def mock_init(update_callback: Callable[[int | None], None] | None):
+        def mock_init(
+            update_callback: Callable[[int | None], None] | None,
+            monitor_height: bool = True,
+        ):
             mock_desk.trigger_update_callback = update_callback
             return mock_desk
 
         desk_patched.side_effect = mock_init
 
-        async def mock_connect(ble_device, monitor_height: bool = True):
+        async def mock_connect(ble_device):
             mock_desk.is_connected = True
+            mock_desk.trigger_update_callback(None)
+
+        async def mock_disconnect():
+            mock_desk.is_connected = False
+            mock_desk.trigger_update_callback(None)
 
         async def mock_move_to(height: float):
             mock_desk.height_percent = height
@@ -38,12 +50,13 @@ async def mock_move_down():
             await mock_move_to(0)
 
         mock_desk.connect = AsyncMock(side_effect=mock_connect)
-        mock_desk.disconnect = AsyncMock()
+        mock_desk.disconnect = AsyncMock(side_effect=mock_disconnect)
         mock_desk.move_to = AsyncMock(side_effect=mock_move_to)
         mock_desk.move_up = AsyncMock(side_effect=mock_move_up)
         mock_desk.move_down = AsyncMock(side_effect=mock_move_down)
         mock_desk.stop = AsyncMock()
         mock_desk.height_percent = 60
         mock_desk.is_moving = False
+        mock_desk.address = "AA:BB:CC:DD:EE:FF"
 
         yield mock_desk

tests/components/idasen_desk/test_config_flow.py

@@ -1,8 +1,8 @@
 """Test the IKEA Idasen Desk config flow."""
-from unittest.mock import patch
+from unittest.mock import ANY, patch
 
-from bleak import BleakError
-from idasen_ha import AuthFailedError
+from bleak.exc import BleakError
+from idasen_ha.errors import AuthFailedError
 import pytest
 
 from homeassistant import config_entries
@@ -260,7 +260,9 @@ async def test_bluetooth_step_success(hass: HomeAssistant) -> None:
     assert result["step_id"] == "user"
     assert result["errors"] == {}
 
-    with patch("homeassistant.components.idasen_desk.config_flow.Desk.connect"), patch(
+    with patch(
+        "homeassistant.components.idasen_desk.config_flow.Desk.connect"
+    ) as desk_connect, patch(
         "homeassistant.components.idasen_desk.config_flow.Desk.disconnect"
     ), patch(
         "homeassistant.components.idasen_desk.async_setup_entry",
@@ -281,3 +283,4 @@ async def test_bluetooth_step_success(hass: HomeAssistant) -> None:
     }
     assert result2["result"].unique_id == IDASEN_DISCOVERY_INFO.address
     assert len(mock_setup_entry.mock_calls) == 1
+    desk_connect.assert_called_with(ANY, auto_reconnect=False)

tests/components/idasen_desk/test_init.py

@@ -1,7 +1,9 @@
 """Test the IKEA Idasen Desk init."""
+from unittest import mock
 from unittest.mock import AsyncMock, MagicMock
 
-from bleak import BleakError
+from bleak.exc import BleakError
+from idasen_ha.errors import AuthFailedError
 import pytest
 
 from homeassistant.components.idasen_desk.const import DOMAIN
@@ -28,7 +30,7 @@ async def test_setup_and_shutdown(
     mock_desk_api.disconnect.assert_called_once()
 
 
-@pytest.mark.parametrize("exception", [TimeoutError(), BleakError()])
+@pytest.mark.parametrize("exception", [AuthFailedError(), TimeoutError(), BleakError()])
 async def test_setup_connect_exception(
     hass: HomeAssistant, mock_desk_api: MagicMock, exception: Exception
 ) -> None:
@@ -39,6 +41,17 @@ async def test_setup_connect_exception(
     assert entry.state is ConfigEntryState.SETUP_RETRY
 
 
+async def test_no_ble_device(hass: HomeAssistant, mock_desk_api: MagicMock) -> None:
+    """Test setup with no BLEDevice from address."""
+    with mock.patch(
+        "homeassistant.components.idasen_desk.bluetooth.async_ble_device_from_address",
+        return_value=None,
+    ):
+        entry = await init_integration(hass)
+        assert len(hass.config_entries.async_entries(DOMAIN)) == 1
+        assert entry.state is ConfigEntryState.SETUP_RETRY
+
+
 async def test_unload_entry(hass: HomeAssistant, mock_desk_api: MagicMock) -> None:
     """Test successful unload of entry."""
     entry = await init_integration(hass)