From 606b76c6815ff2faac78bcd6e825ea91b190c4c7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ab=C3=ADlio=20Costa?= Date: Wed, 18 Oct 2023 23:58:31 +0100 Subject: [PATCH] Add better connection management for Idasen Desk (#102135) --- .../components/idasen_desk/__init__.py | 78 +++++++++++++++---- .../components/idasen_desk/config_flow.py | 7 +- homeassistant/components/idasen_desk/cover.py | 19 ++--- .../components/idasen_desk/manifest.json | 2 +- requirements_all.txt | 2 +- requirements_test_all.txt | 2 +- tests/components/idasen_desk/conftest.py | 19 ++++- .../idasen_desk/test_config_flow.py | 11 ++- tests/components/idasen_desk/test_init.py | 17 +++- 9 files changed, 111 insertions(+), 46 deletions(-) diff --git a/homeassistant/components/idasen_desk/__init__.py b/homeassistant/components/idasen_desk/__init__.py index 04b3ef22e1b..27e7e872fd5 100644 --- a/homeassistant/components/idasen_desk/__init__.py +++ b/homeassistant/components/idasen_desk/__init__.py @@ -4,8 +4,9 @@ from __future__ import annotations 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 @@ from homeassistant.const import ( 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,41 +29,84 @@ PLATFORMS: list[Platform] = [Platform.COVER] _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) @@ -70,7 +114,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: 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 diff --git a/homeassistant/components/idasen_desk/config_flow.py b/homeassistant/components/idasen_desk/config_flow.py index 92f5a836751..caa8d866fc3 100644 --- a/homeassistant/components/idasen_desk/config_flow.py +++ b/homeassistant/components/idasen_desk/config_flow.py @@ -6,7 +6,8 @@ from typing import Any 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 @@ class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN): ) 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" diff --git a/homeassistant/components/idasen_desk/cover.py b/homeassistant/components/idasen_desk/cover.py index c1d1bb48fd8..94f1b4a8cda 100644 --- a/homeassistant/components/idasen_desk/cover.py +++ b/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.const import ATTR_NAME 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 diff --git a/homeassistant/components/idasen_desk/manifest.json b/homeassistant/components/idasen_desk/manifest.json index cdb06cf907d..ed941f4f87d 100644 --- a/homeassistant/components/idasen_desk/manifest.json +++ b/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"] } diff --git a/requirements_all.txt b/requirements_all.txt index cda3a1ca0ec..26ecd839da9 100644 --- a/requirements_all.txt +++ b/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 diff --git a/requirements_test_all.txt b/requirements_test_all.txt index b5d18465769..1f906afb7ef 100644 --- a/requirements_test_all.txt +++ b/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 diff --git a/tests/components/idasen_desk/conftest.py b/tests/components/idasen_desk/conftest.py index 736bc6346ce..d6c2ba5ad6b 100644 --- a/tests/components/idasen_desk/conftest.py +++ b/tests/components/idasen_desk/conftest.py @@ -10,6 +10,10 @@ import pytest @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 @@ def mock_desk_api(): 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 diff --git a/tests/components/idasen_desk/test_config_flow.py b/tests/components/idasen_desk/test_config_flow.py index 223ecc55e28..ca585c65e4d 100644 --- a/tests/components/idasen_desk/test_config_flow.py +++ b/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) diff --git a/tests/components/idasen_desk/test_init.py b/tests/components/idasen_desk/test_init.py index e596f0fe000..cc8daaf98ea 100644 --- a/tests/components/idasen_desk/test_init.py +++ b/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)