Split timeout in lutron_caseta to increase configure timeout (#138875)

homeassistant/components/lutron_caseta/__init__.py

@@ -3,7 +3,6 @@
 from __future__ import annotations
 
 import asyncio
-import contextlib
 from itertools import chain
 import logging
 import ssl
@@ -37,11 +36,12 @@ from .const import (
     ATTR_SERIAL,
     ATTR_TYPE,
     BRIDGE_DEVICE_ID,
-    BRIDGE_TIMEOUT,
     CONF_CA_CERTS,
     CONF_CERTFILE,
     CONF_KEYFILE,
     CONF_SUBTYPE,
+    CONFIGURE_TIMEOUT,
+    CONNECT_TIMEOUT,
     DOMAIN,
     LUTRON_CASETA_BUTTON_EVENT,
     MANUFACTURER,
@@ -161,28 +161,40 @@ async def async_setup_entry(
     keyfile = hass.config.path(entry.data[CONF_KEYFILE])
     certfile = hass.config.path(entry.data[CONF_CERTFILE])
     ca_certs = hass.config.path(entry.data[CONF_CA_CERTS])
-    bridge = None
+    connected_future: asyncio.Future[None] = hass.loop.create_future()
+
+    def _on_connect() -> None:
+        nonlocal connected_future
+        if not connected_future.done():
+            connected_future.set_result(None)
 
     try:
         bridge = Smartbridge.create_tls(
-            hostname=host, keyfile=keyfile, certfile=certfile, ca_certs=ca_certs
+            hostname=host,
+            keyfile=keyfile,
+            certfile=certfile,
+            ca_certs=ca_certs,
+            on_connect_callback=_on_connect,
         )
     except ssl.SSLError:
         _LOGGER.error("Invalid certificate used to connect to bridge at %s", host)
         return False
 
-    timed_out = True
-    with contextlib.suppress(TimeoutError):
-        async with asyncio.timeout(BRIDGE_TIMEOUT):
-            await bridge.connect()
-            timed_out = False
+    connect_task = hass.async_create_task(bridge.connect())
+    for future, name, timeout in (
+        (connected_future, "connect", CONNECT_TIMEOUT),
+        (connect_task, "configure", CONFIGURE_TIMEOUT),
+    ):
+        try:
+            async with asyncio.timeout(timeout):
+                await future
+        except TimeoutError as ex:
+            connect_task.cancel()
+            await bridge.close()
+            raise ConfigEntryNotReady(f"Timed out on {name} for {host}") from ex
 
-    if timed_out or not bridge.is_connected():
-        await bridge.close()
-        if timed_out:
-            raise ConfigEntryNotReady(f"Timed out while trying to connect to {host}")
-        if not bridge.is_connected():
-            raise ConfigEntryNotReady(f"Cannot connect to {host}")
+    if not bridge.is_connected():
+        raise ConfigEntryNotReady(f"Connection failed to {host}")
 
     _LOGGER.debug("Connected to Lutron Caseta bridge via LEAP at %s", host)
     await _async_migrate_unique_ids(hass, entry)

homeassistant/components/lutron_caseta/config_flow.py

@@ -20,10 +20,11 @@ from homeassistant.helpers.service_info.zeroconf import ZeroconfServiceInfo
 from .const import (
     ABORT_REASON_CANNOT_CONNECT,
     BRIDGE_DEVICE_ID,
-    BRIDGE_TIMEOUT,
     CONF_CA_CERTS,
     CONF_CERTFILE,
     CONF_KEYFILE,
+    CONFIGURE_TIMEOUT,
+    CONNECT_TIMEOUT,
     DOMAIN,
     ERROR_CANNOT_CONNECT,
     STEP_IMPORT_FAILED,
@@ -232,7 +233,7 @@ class LutronCasetaFlowHandler(ConfigFlow, domain=DOMAIN):
             return None
 
         try:
-            async with asyncio.timeout(BRIDGE_TIMEOUT):
+            async with asyncio.timeout(CONNECT_TIMEOUT + CONFIGURE_TIMEOUT):
                 await bridge.connect()
         except TimeoutError:
             _LOGGER.error(

homeassistant/components/lutron_caseta/const.py

@@ -34,7 +34,8 @@ ACTION_RELEASE = "release"
 
 CONF_SUBTYPE = "subtype"
 
-BRIDGE_TIMEOUT = 35
+CONNECT_TIMEOUT = 9
+CONFIGURE_TIMEOUT = 50
 
 UNASSIGNED_AREA = "Unassigned"
 

tests/components/lutron_caseta/__init__.py

@@ -1,5 +1,8 @@
 """Tests for the Lutron Caseta integration."""
 
+import asyncio
+from collections.abc import Callable
+from typing import Any
 from unittest.mock import patch
 
 from homeassistant.components.lutron_caseta import DOMAIN
@@ -84,25 +87,12 @@ _LEAP_DEVICE_TYPES = {
 }
 
 
-async def async_setup_integration(hass: HomeAssistant, mock_bridge) -> MockConfigEntry:
-    """Set up a mock bridge."""
-    mock_entry = MockConfigEntry(domain=DOMAIN, data=ENTRY_MOCK_DATA)
-    mock_entry.add_to_hass(hass)
-
-    with patch(
-        "homeassistant.components.lutron_caseta.Smartbridge.create_tls"
-    ) as create_tls:
-        create_tls.return_value = mock_bridge(can_connect=True)
-        await hass.config_entries.async_setup(mock_entry.entry_id)
-        await hass.async_block_till_done()
-    return mock_entry
-
-
 class MockBridge:
     """Mock Lutron bridge that emulates configured connected status."""
 
-    def __init__(self, can_connect=True) -> None:
+    def __init__(self, can_connect=True, timeout_on_connect=False) -> None:
         """Initialize MockBridge instance with configured mock connectivity."""
+        self.timeout_on_connect = timeout_on_connect
         self.can_connect = can_connect
         self.is_currently_connected = False
         self.areas = self.load_areas()
@@ -113,6 +103,8 @@ class MockBridge:
 
     async def connect(self):
         """Connect the mock bridge."""
+        if self.timeout_on_connect:
+            await asyncio.Event().wait()  # wait forever
         if self.can_connect:
             self.is_currently_connected = True
 
@@ -320,3 +312,43 @@ class MockBridge:
     async def close(self):
         """Close the mock bridge connection."""
         self.is_currently_connected = False
+
+
+def make_mock_entry() -> MockConfigEntry:
+    """Create a mock config entry."""
+    return MockConfigEntry(domain=DOMAIN, data=ENTRY_MOCK_DATA)
+
+
+async def async_setup_integration(
+    hass: HomeAssistant,
+    mock_bridge: MockBridge,
+    config_entry_id: str | None = None,
+    can_connect: bool = True,
+    timeout_during_connect: bool = False,
+    timeout_during_configure: bool = False,
+) -> MockConfigEntry:
+    """Set up a mock bridge."""
+    if config_entry_id is None:
+        mock_entry = make_mock_entry()
+        mock_entry.add_to_hass(hass)
+        config_entry_id = mock_entry.entry_id
+    else:
+        mock_entry = hass.config_entries.async_get_entry(config_entry_id)
+
+    def create_tls_factory(
+        *args: Any, on_connect_callback: Callable[[], None], **kwargs: Any
+    ) -> None:
+        """Return a mock bridge."""
+        if not timeout_during_connect:
+            on_connect_callback()
+        return mock_bridge(
+            can_connect=can_connect, timeout_on_connect=timeout_during_configure
+        )
+
+    with patch(
+        "homeassistant.components.lutron_caseta.Smartbridge.create_tls",
+        create_tls_factory,
+    ):
+        await hass.config_entries.async_setup(config_entry_id)
+        await hass.async_block_till_done()
+    return mock_entry

tests/components/lutron_caseta/test_device_trigger.py

@@ -1,7 +1,5 @@
 """The tests for Lutron Caséta device triggers."""
 
-from unittest.mock import patch
-
 import pytest
 from pytest_unordered import unordered
 
@@ -37,7 +35,7 @@ from homeassistant.core import HomeAssistant, ServiceCall
 from homeassistant.helpers import device_registry as dr
 from homeassistant.setup import async_setup_component
 
-from . import MockBridge
+from . import MockBridge, async_setup_integration
 
 from tests.common import MockConfigEntry, async_get_device_automations
 
@@ -112,12 +110,7 @@ async def _async_setup_lutron_with_picos(hass: HomeAssistant) -> str:
     )
     config_entry.add_to_hass(hass)
 
-    with patch(
-        "homeassistant.components.lutron_caseta.Smartbridge.create_tls",
-        return_value=MockBridge(can_connect=True),
-    ):
-        await hass.config_entries.async_setup(config_entry.entry_id)
-        await hass.async_block_till_done()
+    await async_setup_integration(hass, MockBridge, config_entry.entry_id)
 
     return config_entry.entry_id
 
@@ -487,9 +480,7 @@ async def test_if_fires_on_button_event_late_setup(
         },
     )
 
-    with patch("homeassistant.components.lutron_caseta.Smartbridge.create_tls"):
-        await hass.config_entries.async_setup(config_entry_id)
-        await hass.async_block_till_done()
+    await async_setup_integration(hass, MockBridge, config_entry_id)
 
     message = {
         ATTR_SERIAL: device.get("serial"),

tests/components/lutron_caseta/test_diagnostics.py

@@ -1,6 +1,6 @@
 """Test the Lutron Caseta diagnostics."""
 
-from unittest.mock import ANY, patch
+from unittest.mock import ANY
 
 from homeassistant.components.lutron_caseta import DOMAIN
 from homeassistant.components.lutron_caseta.const import (
@@ -11,7 +11,7 @@ from homeassistant.components.lutron_caseta.const import (
 from homeassistant.const import CONF_HOST
 from homeassistant.core import HomeAssistant
 
-from . import MockBridge
+from . import MockBridge, async_setup_integration
 
 from tests.common import MockConfigEntry
 from tests.components.diagnostics import get_diagnostics_for_config_entry
@@ -34,12 +34,7 @@ async def test_diagnostics(
     )
     config_entry.add_to_hass(hass)
 
-    with patch(
-        "homeassistant.components.lutron_caseta.Smartbridge.create_tls",
-        return_value=MockBridge(can_connect=True),
-    ):
-        await hass.config_entries.async_setup(config_entry.entry_id)
-        await hass.async_block_till_done()
+    await async_setup_integration(hass, MockBridge, config_entry.entry_id)
 
     diag = await get_diagnostics_for_config_entry(hass, hass_client, config_entry)
     assert diag == {

tests/components/lutron_caseta/test_init.py

@@ -0,0 +1,54 @@
+"""Tests for the Lutron Caseta integration."""
+
+from unittest.mock import patch
+
+import pytest
+
+from homeassistant.components import lutron_caseta
+from homeassistant.config_entries import ConfigEntryState
+from homeassistant.core import HomeAssistant
+
+from . import MockBridge, async_setup_integration, make_mock_entry
+
+
+@pytest.mark.parametrize(
+    ("constant", "message", "timeout_during_connect", "timeout_during_configure"),
+    [
+        ("CONNECT_TIMEOUT", "Timed out on connect", True, False),
+        ("CONFIGURE_TIMEOUT", "Timed out on configure", False, True),
+    ],
+)
+async def test_timeout_during_setup(
+    hass: HomeAssistant,
+    caplog: pytest.LogCaptureFixture,
+    constant: str,
+    message: str,
+    timeout_during_connect: bool,
+    timeout_during_configure: bool,
+) -> None:
+    """Test a timeout during setup."""
+    mock_entry = make_mock_entry()
+    mock_entry.add_to_hass(hass)
+    with patch.object(lutron_caseta, constant, 0.001):
+        await async_setup_integration(
+            hass,
+            MockBridge,
+            config_entry_id=mock_entry.entry_id,
+            timeout_during_connect=timeout_during_connect,
+            timeout_during_configure=timeout_during_configure,
+        )
+    assert mock_entry.state is ConfigEntryState.SETUP_RETRY
+    assert f"{message} for 1.1.1.1" in caplog.text
+
+
+async def test_cannot_connect(
+    hass: HomeAssistant, caplog: pytest.LogCaptureFixture
+) -> None:
+    """Test failing to connect."""
+    mock_entry = make_mock_entry()
+    mock_entry.add_to_hass(hass)
+    await async_setup_integration(
+        hass, MockBridge, config_entry_id=mock_entry.entry_id, can_connect=False
+    )
+    assert mock_entry.state is ConfigEntryState.SETUP_RETRY
+    assert "Connection failed to 1.1.1.1" in caplog.text

tests/components/lutron_caseta/test_logbook.py

@@ -1,7 +1,5 @@
 """The tests for lutron caseta logbook."""
 
-from unittest.mock import patch
-
 from homeassistant.components.lutron_caseta.const import (
     ATTR_ACTION,
     ATTR_AREA_NAME,
@@ -43,13 +41,7 @@ async def test_humanify_lutron_caseta_button_event(hass: HomeAssistant) -> None:
         unique_id="abc",
     )
     config_entry.add_to_hass(hass)
-
-    with patch(
-        "homeassistant.components.lutron_caseta.Smartbridge.create_tls",
-        return_value=MockBridge(can_connect=True),
-    ):
-        await hass.config_entries.async_setup(config_entry.entry_id)
-        await hass.async_block_till_done()
+    await async_setup_integration(hass, MockBridge, config_entry.entry_id)
 
     await hass.async_block_till_done()
 
@@ -104,15 +96,10 @@ async def test_humanify_lutron_caseta_button_event_integration_not_loaded(
     )
     config_entry.add_to_hass(hass)
 
-    with patch(
-        "homeassistant.components.lutron_caseta.Smartbridge.create_tls",
-        return_value=MockBridge(can_connect=True),
-    ):
-        await hass.config_entries.async_setup(config_entry.entry_id)
-        await hass.async_block_till_done()
+    await async_setup_integration(hass, MockBridge, config_entry.entry_id)
 
-        await hass.config_entries.async_unload(config_entry.entry_id)
-        await hass.async_block_till_done()
+    await hass.config_entries.async_unload(config_entry.entry_id)
+    await hass.async_block_till_done()
 
     for device in device_registry.devices.values():
         if device.config_entries == {config_entry.entry_id}: