Automatically keep connection state (#28)

This moves the logic to keep the device connected when the connection
drops (and vice-versa) from HA to this lib.

idasen_ha/__init__.py

@@ -6,7 +6,6 @@
 import logging
 from typing import Callable
 
-from bleak import BleakClient
 from bleak.backends.device import BLEDevice
 from idasen import IdasenDesk
 
@@ -24,7 +23,6 @@ def __init__(
         monitor_height: bool = True,
     ) -> None:
         """Initialize the wrapper."""
-        self._idasen_desk: IdasenDesk | None = None
         self._ble_device: BLEDevice | None = None
         self._connection_manager: ConnectionManager | None = None
         self._height: float | None = None
@@ -51,12 +49,9 @@ async def connect(self, ble_device: BLEDevice, retry: bool = True) -> None:
             _LOGGER.debug("Initializing idasen desk")
             await self.disconnect()
             self._ble_device = ble_device
-            self._idasen_desk = self._create_idasen_desk(self._ble_device)
-            self._connection_manager = self._create_connection_manager(
-                self._idasen_desk
-            )
+            self._connection_manager = self._create_connection_manager(self._ble_device)
 
-        await self._connection_manager.connect(retry=retry)
+        await self._connection_manager.connect(retry)
 
     async def disconnect(self) -> None:
         """Disconnect from the desk."""
@@ -141,17 +136,13 @@ def is_connected(self) -> bool:
         # will always be False.
         return bool(self._idasen_desk.is_connected)
 
-    def _create_idasen_desk(self, ble_device: BLEDevice) -> IdasenDesk:
-        def disconnect_callback(client: BleakClient) -> None:
-            """Handle bluetooth disconnection."""
-            _LOGGER.debug("Disconnect callback called")
-            self._update_callback(self.height_percent)
-
-        return IdasenDesk(
-            ble_device, exit_on_fail=False, disconnected_callback=disconnect_callback
-        )
+    @property
+    def _idasen_desk(self) -> IdasenDesk | None:
+        if self._connection_manager is None:
+            return None
+        return self._connection_manager.idasen_desk
 
-    def _create_connection_manager(self, desk: IdasenDesk) -> ConnectionManager:
+    def _create_connection_manager(self, ble_device: BLEDevice) -> ConnectionManager:
         async def connect_callback() -> None:
             _LOGGER.debug("Connect callback called")
             if self._idasen_desk is None:
@@ -165,6 +156,7 @@ async def connect_callback() -> None:
             self._update_callback(self.height_percent)
 
         return ConnectionManager(
-            desk,
+            ble_device,
             connect_callback=connect_callback,
+            disconnect_callback=lambda: self._update_callback(self.height_percent),
         )

idasen_ha/connection_manager.py

@@ -5,6 +5,7 @@
 import logging
 from typing import Callable
 
+from bleak.backends.device import BLEDevice
 from bleak.exc import BleakDBusError, BleakError
 from idasen import IdasenDesk
 
@@ -17,27 +18,44 @@ class ConnectionManager:
     """Handles connecting to the desk. Optionally keeps retrying to connect until it succeeds."""
 
     def __init__(
-        self, desk: IdasenDesk, connect_callback: Callable[[], Awaitable[None]]
+        self,
+        ble_device: BLEDevice,
+        connect_callback: Callable[[], Awaitable[None]],
+        disconnect_callback: Callable[[], None],
     ):
         """Init ConnectionManager."""
-        self._idasen_desk = desk
-        self._connect_callback = connect_callback
-
+        self._keep_connected: bool = False
         self._connecting: bool = False
         self._retry_pending: bool = False
-        self._retry: bool = False
 
-    async def connect(self, retry: bool = True) -> None:
+        self._idasen_desk: IdasenDesk = self._create_idasen_desk(ble_device)
+
+        self._connect_callback = connect_callback
+        self._disconnect_callback = disconnect_callback
+
+    @property
+    def idasen_desk(self) -> IdasenDesk:
+        """The IdasenDesk instance."""
+        return self._idasen_desk
+
+    async def connect(self, retry: bool) -> None:
         """Perform the bluetooth connection to the desk."""
-        self._retry = retry
-        await self._connect(retry=retry)
+        self._keep_connected = True
+        await self._connect(retry)
 
     async def disconnect(self):
         """Stop the connection manager retry task."""
-        self._retry = False
+        self._keep_connected = False
         if self._idasen_desk.is_connected:
             await self._idasen_desk.disconnect()
 
+    def _create_idasen_desk(self, ble_device: BLEDevice) -> IdasenDesk:
+        return IdasenDesk(
+            ble_device,
+            exit_on_fail=False,
+            disconnected_callback=lambda bledevice: self._handle_disconnect(),
+        )
+
     async def _connect(self, retry: bool) -> None:
         if self._connecting:
             _LOGGER.info("Connection already in progress.")
@@ -49,7 +67,7 @@ async def _connect(self, retry: bool) -> None:
                 _LOGGER.info("Connecting...")
                 await self._idasen_desk.connect()
             except (TimeoutError, BleakError) as ex:
-                _LOGGER.warning("Connect failed")
+                _LOGGER.exception("Connect failed")
                 if retry:
                     self._schedule_reconnect()
                     return
@@ -65,40 +83,64 @@ async def _connect(self, retry: bool) -> None:
                     raise AuthFailedError() from ex
                 raise ex
             except Exception as ex:
-                _LOGGER.warning("Pair failed")
+                _LOGGER.exception("Pair failed")
                 await self._idasen_desk.disconnect()
                 if retry:
                     self._schedule_reconnect()
                     return
                 else:
                     raise ex
 
-            _LOGGER.info("Connected!")
-            self._retry = False
-            await self._connect_callback()
+            await self._handle_connect()
         finally:
             self._connecting = False
 
     def _schedule_reconnect(self):
         RECONNECT_INTERVAL_SEC = 30
-        _LOGGER.info("Will try to connect in %ds", RECONNECT_INTERVAL_SEC)
+        _LOGGER.info("Will try to reconnect in %ds", RECONNECT_INTERVAL_SEC)
 
         if self._retry_pending:
             _LOGGER.warning("There is already a reconnect task pending")
             return
+        self._retry_pending = True
 
         async def _reconnect():
-            self._retry_pending = True
             await asyncio.sleep(RECONNECT_INTERVAL_SEC)
             self._retry_pending = False
 
-            if self._retry is False:
+            if not self._keep_connected:
                 _LOGGER.debug(
-                    "Retrying is disalbed (this could be on an older instance for an older BLEDevice)"
+                    "Reconnect aborted since it should not be connected"
+                    "(this could be on an older instance for an older BLEDevice)"
                 )
                 return
 
-            _LOGGER.debug("Retrying to connect now")
-            await self._connect(retry=True)
+            if self.idasen_desk.is_connected:
+                _LOGGER.debug("Already connected")
+                return
+
+            _LOGGER.debug("Reconnecting now")
+            await self._connect(True)
 
         asyncio.get_event_loop().create_task(_reconnect())
+
+    async def _handle_connect(self) -> None:
+        _LOGGER.debug("Connected")
+        await self._connect_callback()
+        if not self._keep_connected:
+            _LOGGER.info("Disconnecting since it should not be connected")
+            asyncio.get_event_loop().create_task(self.disconnect())
+
+    def _handle_disconnect(self) -> None:
+        """Handle bluetooth disconnection."""
+        _LOGGER.debug("Disconnected")
+        if self._connecting:
+            _LOGGER.debug(
+                "Disconnected during connection process. No callback triggered."
+            )
+            return
+
+        self._disconnect_callback()
+        if self._keep_connected and not self._retry_pending:
+            _LOGGER.info("Reconnecting since it should not be disconnected")
+            asyncio.get_event_loop().create_task(self.connect(True))

tests/__init__.py

@@ -3,7 +3,7 @@
 from bleak.backends.device import BLEDevice
 from idasen import IdasenDesk
 
-FAKE_BLE_DEVICE = BLEDevice("AA:BB:CC:DD:EE:FF", None, None, 0)
+FAKE_BLE_DEVICE = BLEDevice("AA:BB:CC:DD:EE:FF", None, {"path": ""}, 0)
 
 
 def height_percent_to_meters(percent: float):

tests/conftest.py

@@ -14,7 +14,9 @@
 async def mock_idasen_desk():
     """Test height monitoring."""
 
-    with mock.patch("idasen_ha.IdasenDesk", autospec=True) as patched_idasen_desk:
+    with mock.patch(
+        "idasen_ha.connection_manager.IdasenDesk", autospec=True
+    ) as patched_idasen_desk:
         patched_idasen_desk.MIN_HEIGHT = IdasenDesk.MIN_HEIGHT
         patched_idasen_desk.MAX_HEIGHT = IdasenDesk.MAX_HEIGHT
 

tests/test_connection_management.py

@@ -56,7 +56,9 @@ async def connect_side_effect():
 async def test_double_connect_call_with_different_bledevice():
     """Test connect being called again with a new BLEDevice, while still connecting."""
 
-    with mock.patch("idasen_ha.IdasenDesk", autospec=True) as patched_idasen_desk:
+    with mock.patch(
+        "idasen_ha.connection_manager.IdasenDesk", autospec=True
+    ) as patched_idasen_desk:
         mock_idasen_desk = patched_idasen_desk.return_value
 
         async def connect_side_effect():
@@ -136,7 +138,7 @@ async def test_disconnect_on_pair_failure(
         await desk.connect(FAKE_BLE_DEVICE, retry=False)
     assert not desk.is_connected
     mock_idasen_desk.disconnect.assert_called()
-    assert update_callback.call_count == 1
+    assert update_callback.call_count == 0
 
 
 @mock.patch("idasen_ha.connection_manager.asyncio.sleep")
@@ -203,3 +205,24 @@ async def sleep_handler(delay):
 
     await retry_maxed_future
     assert mock_idasen_desk.connect.call_count == TEST_RETRIES_MAX + 2
+
+
+async def test_reconnect_on_connection_drop(mock_idasen_desk: MagicMock):
+    """Test reconnection when the connection drops."""
+    update_callback = Mock()
+    desk = Desk(update_callback, True)
+
+    await desk.connect(FAKE_BLE_DEVICE)
+    assert desk.is_connected
+    assert update_callback.call_count == 1
+
+    mock_idasen_desk.reset_mock()
+    await mock_idasen_desk.disconnect()
+    assert not desk.is_connected
+    assert update_callback.call_count == 2
+
+    await asyncio.sleep(0)
+    assert desk.is_connected
+    mock_idasen_desk.connect.assert_awaited()
+    mock_idasen_desk.pair.assert_called()
+    assert update_callback.call_count == 3