Improve Droplet connection handling (#184092)

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,

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"

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

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,
         ),
     ):

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)