From ec8b79d1d49719df0c2ec414ed000f1e34930fdc Mon Sep 17 00:00:00 2001 From: Sarah Seidman Date: Mon, 5 Oct 2026 07:58:18 -0400 Subject: [PATCH] Improve Droplet connection handling (#184092) --- homeassistant/components/droplet/config_flow.py | 16 ++++++++++++---- homeassistant/components/droplet/const.py | 3 ++- homeassistant/components/droplet/coordinator.py | 12 +++++------- tests/components/droplet/conftest.py | 4 ++-- tests/components/droplet/test_config_flow.py | 3 +++ 5 files changed, 24 insertions(+), 14 deletions(-) diff --git a/homeassistant/components/droplet/config_flow.py b/homeassistant/components/droplet/config_flow.py index 1f971e522656..46c4f46d2ec2 100644 --- a/homeassistant/components/droplet/config_flow.py +++ b/homeassistant/components/droplet/config_flow.py @@ -57,7 +57,11 @@ class DropletConfigFlow(ConfigFlow, domain=DOMAIN): # Test if we can connect before returning session = async_get_clientsession(self.hass) code = normalize_pairing_code(user_input[CONF_CODE]) - if await self._droplet_discovery.try_connect(session, code): + try: + connected = await self._droplet_discovery.try_connect(session, code) + finally: + await self._droplet_discovery.close() + if connected: device_data = { CONF_IP_ADDRESS: self._droplet_discovery.host, CONF_PORT: self._droplet_discovery.port, @@ -95,9 +99,13 @@ class DropletConfigFlow(ConfigFlow, domain=DOMAIN): ) session = async_get_clientsession(self.hass) code = normalize_pairing_code(user_input[CONF_CODE]) - if await self._droplet_discovery.try_connect(session, code) and ( - device_id := await self._droplet_discovery.get_device_id() - ): + device_id = "" + try: + if await self._droplet_discovery.try_connect(session, code): + device_id = await self._droplet_discovery.get_device_id() + finally: + await self._droplet_discovery.close() + if device_id: device_data = { CONF_IP_ADDRESS: self._droplet_discovery.host, CONF_PORT: self._droplet_discovery.port, diff --git a/homeassistant/components/droplet/const.py b/homeassistant/components/droplet/const.py index fa9df48a02d1..1f55507ce8eb 100644 --- a/homeassistant/components/droplet/const.py +++ b/homeassistant/components/droplet/const.py @@ -1,6 +1,7 @@ """Constants for the droplet integration.""" -CONNECT_DELAY = 5 +CONNECT_TIMEOUT = 15 +RECONNECT_DELAY = 5 DOMAIN = "droplet" DEVICE_NAME = "Droplet" diff --git a/homeassistant/components/droplet/coordinator.py b/homeassistant/components/droplet/coordinator.py index 1ee324536453..31a0d41e8ab3 100644 --- a/homeassistant/components/droplet/coordinator.py +++ b/homeassistant/components/droplet/coordinator.py @@ -14,9 +14,7 @@ from homeassistant.exceptions import ConfigEntryNotReady from homeassistant.helpers.aiohttp_client import async_get_clientsession from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed -from .const import CONNECT_DELAY, DOMAIN - -VERSION_TIMEOUT = 5 +from .const import CONNECT_TIMEOUT, DOMAIN, RECONNECT_DELAY _LOGGER = logging.getLogger(__name__) @@ -50,8 +48,8 @@ class DropletDataCoordinator(DataUpdateCoordinator[None]): if not await self.setup(): raise ConfigEntryNotReady("Device is offline") - # Droplet should send its metadata within 5 seconds - end = time.time() + VERSION_TIMEOUT + # Droplet should send its metadata shortly after connecting + end = time.time() + CONNECT_TIMEOUT while not self.droplet.version_info_available(): await asyncio.sleep(TIMEOUT) if time.time() > end: @@ -70,10 +68,10 @@ class DropletDataCoordinator(DataUpdateCoordinator[None]): self.config_entry.async_on_unload(self.droplet.stop_listening) self.config_entry.async_create_background_task( self.hass, - self.droplet.listen_forever(CONNECT_DELAY, self.async_set_updated_data), + self.droplet.listen_forever(RECONNECT_DELAY, self.async_set_updated_data), "droplet-listen", ) - end = time.time() + CONNECT_DELAY + end = time.time() + CONNECT_TIMEOUT while time.time() < end: if self.droplet.connected: return True diff --git a/tests/components/droplet/conftest.py b/tests/components/droplet/conftest.py index f3c2a1aac52a..a3631fe464db 100644 --- a/tests/components/droplet/conftest.py +++ b/tests/components/droplet/conftest.py @@ -65,11 +65,11 @@ def mock_timeout() -> Generator[None]: 0.05, ), patch( - "homeassistant.components.droplet.coordinator.VERSION_TIMEOUT", + "homeassistant.components.droplet.coordinator.CONNECT_TIMEOUT", 0.1, ), patch( - "homeassistant.components.droplet.coordinator.CONNECT_DELAY", + "homeassistant.components.droplet.coordinator.RECONNECT_DELAY", 0.1, ), ): diff --git a/tests/components/droplet/test_config_flow.py b/tests/components/droplet/test_config_flow.py index f706a94eb748..9f979767e5c4 100644 --- a/tests/components/droplet/test_config_flow.py +++ b/tests/components/droplet/test_config_flow.py @@ -65,6 +65,7 @@ async def test_user_setup( } assert result.get("context") is not None assert result.get("context", {}).get("unique_id") == MOCK_DEVICE_ID + mock_droplet_discovery.close.assert_awaited_once() @pytest.mark.parametrize( @@ -107,6 +108,7 @@ async def test_user_setup_fail( assert result is not None assert result.get("type") is FlowResultType.FORM assert result.get("errors") == {"base": "cannot_connect"} + mock_droplet_discovery.close.assert_awaited_once() # The user should be able to try again. Maybe the droplet # was disconnected from the network or something @@ -200,6 +202,7 @@ async def test_zeroconf_setup( } assert result.get("context") is not None assert result.get("context", {}).get("unique_id") == MOCK_DEVICE_ID + mock_droplet_discovery.close.assert_awaited_once() @pytest.mark.parametrize("mock_droplet_discovery", ["192.168.1.5"], indirect=True)