Fix System Bridge wait timeout wait condition (#141811)

* Fix System Bridge wait timeout wait condition

* Add DataMissingException as a timeout condition

* Add tests

homeassistant/components/system_bridge/__init__.py

@@ -11,6 +11,7 @@ from systembridgeconnector.exceptions import (
     AuthenticationException,
     ConnectionClosedException,
     ConnectionErrorException,
+    DataMissingException,
 )
 from systembridgeconnector.version import Version
 from systembridgemodels.keyboard_key import KeyboardKey
@@ -184,7 +185,7 @@ async def async_setup_entry(
                 "host": entry.data[CONF_HOST],
             },
         ) from exception
-    except TimeoutError as exception:
+    except (DataMissingException, TimeoutError) as exception:
         raise ConfigEntryNotReady(
             translation_domain=DOMAIN,
             translation_key="timeout",

homeassistant/components/system_bridge/const.py

@@ -18,4 +18,6 @@ MODULES: Final[list[Module]] = [
     Module.SYSTEM,
 ]
 
-DATA_WAIT_TIMEOUT: Final[int] = 10
+DATA_WAIT_TIMEOUT: Final[int] = 20
+
+GET_DATA_WAIT_TIMEOUT: Final[int] = 15

homeassistant/components/system_bridge/coordinator.py

@@ -33,7 +33,7 @@ from homeassistant.exceptions import ConfigEntryAuthFailed
 from homeassistant.helpers.aiohttp_client import async_get_clientsession
 from homeassistant.helpers.update_coordinator import DataUpdateCoordinator
 
-from .const import DOMAIN, MODULES
+from .const import DOMAIN, GET_DATA_WAIT_TIMEOUT, MODULES
 from .data import SystemBridgeData
 
 
@@ -119,7 +119,10 @@ class SystemBridgeDataUpdateCoordinator(DataUpdateCoordinator[SystemBridgeData])
         """Get data from WebSocket."""
         await self.check_websocket_connected()
 
-        modules_data = await self.websocket_client.get_data(GetData(modules=modules))
+        modules_data = await self.websocket_client.get_data(
+            GetData(modules=modules),
+            timeout=GET_DATA_WAIT_TIMEOUT,
+        )
 
         # Merge new data with existing data
         for module in MODULES:

tests/components/system_bridge/test_init.py

@@ -81,3 +81,53 @@ async def test_migration_minor_future_version(hass: HomeAssistant) -> None:
     assert config_entry.minor_version == config_entry_minor_version
     assert config_entry.data == config_entry_data
     assert config_entry.state is ConfigEntryState.LOADED
+
+
+async def test_setup_timeout(hass: HomeAssistant) -> None:
+    """Test setup with timeout error."""
+    config_entry = MockConfigEntry(
+        domain=DOMAIN,
+        unique_id=FIXTURE_UUID,
+        data=FIXTURE_USER_INPUT,
+        version=SystemBridgeConfigFlow.VERSION,
+        minor_version=SystemBridgeConfigFlow.MINOR_VERSION,
+    )
+
+    with patch(
+        "systembridgeconnector.version.Version.check_supported",
+        side_effect=TimeoutError,
+    ):
+        config_entry.add_to_hass(hass)
+        result = await hass.config_entries.async_setup(config_entry.entry_id)
+        await hass.async_block_till_done()
+
+        assert result is False
+        assert config_entry.state is ConfigEntryState.SETUP_RETRY
+
+
+async def test_coordinator_get_data_timeout(hass: HomeAssistant) -> None:
+    """Test coordinator handling timeout during get_data."""
+    config_entry = MockConfigEntry(
+        domain=DOMAIN,
+        unique_id=FIXTURE_UUID,
+        data=FIXTURE_USER_INPUT,
+        version=SystemBridgeConfigFlow.VERSION,
+        minor_version=SystemBridgeConfigFlow.MINOR_VERSION,
+    )
+
+    with (
+        patch(
+            "systembridgeconnector.version.Version.check_supported",
+            return_value=True,
+        ),
+        patch(
+            "homeassistant.components.system_bridge.coordinator.SystemBridgeDataUpdateCoordinator.async_get_data",
+            side_effect=TimeoutError,
+        ),
+    ):
+        config_entry.add_to_hass(hass)
+        result = await hass.config_entries.async_setup(config_entry.entry_id)
+        await hass.async_block_till_done()
+
+        assert result is False
+        assert config_entry.state is ConfigEntryState.SETUP_RETRY