diff --git a/idasen_ha/__init__.py b/idasen_ha/__init__.py index ee36aa8..04d2183 100644 --- a/idasen_ha/__init__.py +++ b/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), ) diff --git a/idasen_ha/connection_manager.py b/idasen_ha/connection_manager.py index 5e103f4..6fb1418 100644 --- a/idasen_ha/connection_manager.py +++ b/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,7 +83,7 @@ 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() @@ -73,32 +91,56 @@ async def _connect(self, retry: bool) -> None: 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)) diff --git a/tests/__init__.py b/tests/__init__.py index 036ac42..f825d9e 100644 --- a/tests/__init__.py +++ b/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): diff --git a/tests/conftest.py b/tests/conftest.py index dde358d..f4d0969 100644 --- a/tests/conftest.py +++ b/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 diff --git a/tests/test_connection_management.py b/tests/test_connection_management.py index f8b087d..2f6aeed 100644 --- a/tests/test_connection_management.py +++ b/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