Fix matter websocket reconnect (#84192)

.coveragerc

@@ -729,7 +729,6 @@ omit =
     homeassistant/components/mastodon/notify.py
     homeassistant/components/matrix/*
     homeassistant/components/matter/__init__.py
-    homeassistant/components/matter/entity.py
     homeassistant/components/meater/__init__.py
     homeassistant/components/meater/const.py
     homeassistant/components/meater/sensor.py

homeassistant/components/matter/__init__.py

@@ -11,10 +11,11 @@ from matter_server.client.exceptions import (
     FailedCommand,
     InvalidServerVersion,
 )
+from matter_server.common.models.error import MatterError
 import voluptuous as vol
 
 from homeassistant.components.hassio import AddonError, AddonManager, AddonState
-from homeassistant.config_entries import ConfigEntry
+from homeassistant.config_entries import ConfigEntry, ConfigEntryState
 from homeassistant.const import CONF_URL, EVENT_HOMEASSISTANT_STOP
 from homeassistant.core import Event, HomeAssistant, ServiceCall, callback
 from homeassistant.exceptions import ConfigEntryNotReady, HomeAssistantError
@@ -32,6 +33,10 @@ from .addon import get_addon_manager
 from .api import async_register_api
 from .const import CONF_INTEGRATION_CREATED_ADDON, CONF_USE_ADDON, DOMAIN, LOGGER
 from .device_platform import DEVICE_PLATFORM
+from .helpers import MatterEntryData, get_matter
+
+CONNECT_TIMEOUT = 10
+LISTEN_READY_TIMEOUT = 30
 
 
 async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
@@ -41,8 +46,9 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
 
     matter_client = MatterClient(entry.data[CONF_URL], async_get_clientsession(hass))
     try:
+        async with async_timeout.timeout(CONNECT_TIMEOUT):
             await matter_client.connect()
-    except CannotConnect as err:
+    except (CannotConnect, asyncio.TimeoutError) as err:
         raise ConfigEntryNotReady("Failed to connect to matter server") from err
     except InvalidServerVersion as err:
         if use_addon:
@@ -60,7 +66,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
         raise ConfigEntryNotReady(f"Invalid server version: {err}") from err
 
     except Exception as err:
-        matter_client.logger.exception("Failed to connect to matter server")
+        LOGGER.exception("Failed to connect to matter server")
         raise ConfigEntryNotReady(
             "Unknown error connecting to the Matter server"
         ) from err
@@ -75,16 +81,17 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
         hass.bus.async_listen_once(EVENT_HOMEASSISTANT_STOP, on_hass_stop)
     )
 
-    # register websocket api
     async_register_api(hass)
 
     # launch the matter client listen task in the background
-    # use the init_ready event to keep track if it did initialize successfully
+    # use the init_ready event to wait until initialization is done
     init_ready = asyncio.Event()
-    listen_task = asyncio.create_task(matter_client.start_listening(init_ready))
+    listen_task = asyncio.create_task(
+        _client_listen(hass, entry, matter_client, init_ready)
+    )
 
     try:
-        async with async_timeout.timeout(30):
+        async with async_timeout.timeout(LISTEN_READY_TIMEOUT):
             await init_ready.wait()
     except asyncio.TimeoutError as err:
         listen_task.cancel()
@@ -94,27 +101,58 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
         hass.data[DOMAIN] = {}
         _async_init_services(hass)
 
-    # we create an intermediate layer (adapter) which keeps track of our nodes
+    # create an intermediate layer (adapter) which keeps track of the nodes
-    # and discovery of platform entities from the node's attributes
+    # and discovery of platform entities from the node attributes
     matter = MatterAdapter(hass, matter_client, entry)
-    hass.data[DOMAIN][entry.entry_id] = matter
+    hass.data[DOMAIN][entry.entry_id] = MatterEntryData(matter, listen_task)
 
-    # forward platform setup to all platforms in the discovery schema
     await hass.config_entries.async_forward_entry_setups(entry, DEVICE_PLATFORM)
+    await matter.setup_nodes()
 
-    # start discovering of node entities as task
+    # If the listen task is already failed, we need to raise ConfigEntryNotReady
-    asyncio.create_task(matter.setup_nodes())
+    if listen_task.done() and (listen_error := listen_task.exception()) is not None:
+        await hass.config_entries.async_unload_platforms(entry, DEVICE_PLATFORM)
+        hass.data[DOMAIN].pop(entry.entry_id)
+        try:
+            await matter_client.disconnect()
+        finally:
+            raise ConfigEntryNotReady(listen_error) from listen_error
 
     return True
 
 
+async def _client_listen(
+    hass: HomeAssistant,
+    entry: ConfigEntry,
+    matter_client: MatterClient,
+    init_ready: asyncio.Event,
+) -> None:
+    """Listen with the client."""
+    try:
+        await matter_client.start_listening(init_ready)
+    except MatterError as err:
+        if entry.state != ConfigEntryState.LOADED:
+            raise
+        LOGGER.error("Failed to listen: %s", err)
+    except Exception as err:  # pylint: disable=broad-except
+        # We need to guard against unknown exceptions to not crash this task.
+        LOGGER.exception("Unexpected exception: %s", err)
+        if entry.state != ConfigEntryState.LOADED:
+            raise
+
+    if not hass.is_stopping:
+        LOGGER.debug("Disconnected from server. Reloading integration")
+        hass.async_create_task(hass.config_entries.async_reload(entry.entry_id))
+
+
 async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Unload a config entry."""
     unload_ok = await hass.config_entries.async_unload_platforms(entry, DEVICE_PLATFORM)
 
     if unload_ok:
-        matter: MatterAdapter = hass.data[DOMAIN].pop(entry.entry_id)
+        matter_entry_data: MatterEntryData = hass.data[DOMAIN].pop(entry.entry_id)
-        await matter.matter_client.disconnect()
+        matter_entry_data.listen_task.cancel()
+        await matter_entry_data.adapter.matter_client.disconnect()
 
     if entry.data.get(CONF_USE_ADDON) and entry.disabled_by:
         addon_manager: AddonManager = get_addon_manager(hass)
@@ -165,26 +203,17 @@ async def async_remove_config_entry_device(
     if not unique_id:
         return True
 
-    matter: MatterAdapter = hass.data[DOMAIN][config_entry.entry_id]
+    matter_entry_data: MatterEntryData = hass.data[DOMAIN][config_entry.entry_id]
+    matter_client = matter_entry_data.adapter.matter_client
 
-    for node in await matter.matter_client.get_nodes():
+    for node in await matter_client.get_nodes():
         if node.unique_id == unique_id:
-            await matter.matter_client.remove_node(node.node_id)
+            await matter_client.remove_node(node.node_id)
             break
 
     return True
 
 
-@callback
-def get_matter(hass: HomeAssistant) -> MatterAdapter:
-    """Return MatterAdapter instance."""
-    # NOTE: This assumes only one Matter connection/fabric can exist.
-    # Shall we support connecting to multiple servers in the client or by config entries?
-    # In case of the config entry we need to fix this.
-    matter: MatterAdapter = next(iter(hass.data[DOMAIN].values()))
-    return matter
-
-
 @callback
 def _async_init_services(hass: HomeAssistant) -> None:
     """Init services."""

homeassistant/components/matter/api.py

@@ -13,7 +13,7 @@ from homeassistant.components.websocket_api import ActiveConnection
 from homeassistant.core import HomeAssistant, callback
 
 from .adapter import MatterAdapter
-from .const import DOMAIN
+from .helpers import get_matter
 
 ID = "id"
 TYPE = "type"
@@ -36,7 +36,7 @@ def async_get_matter_adapter(func: Callable) -> Callable:
         hass: HomeAssistant, connection: ActiveConnection, msg: dict
     ) -> None:
         """Provide the Matter client to the function."""
-        matter: MatterAdapter = next(iter(hass.data[DOMAIN].values()))
+        matter = get_matter(hass)
 
         await func(hass, connection, msg, matter)
 

homeassistant/components/matter/binary_sensor.py

@@ -3,7 +3,6 @@ from __future__ import annotations
 
 from dataclasses import dataclass
 from functools import partial
-from typing import TYPE_CHECKING
 
 from chip.clusters import Objects as clusters
 from matter_server.common.models import device_types
@@ -18,11 +17,8 @@ from homeassistant.const import Platform
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from .const import DOMAIN
 from .entity import MatterEntity, MatterEntityDescriptionBaseClass
-
+from .helpers import get_matter
-if TYPE_CHECKING:
-    from .adapter import MatterAdapter
 
 
 async def async_setup_entry(
@@ -31,7 +27,7 @@ async def async_setup_entry(
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up Matter binary sensor from Config Entry."""
-    matter: MatterAdapter = hass.data[DOMAIN][config_entry.entry_id]
+    matter = get_matter(hass)
     matter.register_platform_handler(Platform.BINARY_SENSOR, async_add_entities)
 
 

homeassistant/components/matter/helpers.py

@@ -0,0 +1,31 @@
+"""Provide integration helpers that are aware of the matter integration."""
+from __future__ import annotations
+
+import asyncio
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+from homeassistant.core import HomeAssistant, callback
+
+from .const import DOMAIN
+
+if TYPE_CHECKING:
+    from .adapter import MatterAdapter
+
+
+@dataclass
+class MatterEntryData:
+    """Hold Matter data for the config entry."""
+
+    adapter: MatterAdapter
+    listen_task: asyncio.Task
+
+
+@callback
+def get_matter(hass: HomeAssistant) -> MatterAdapter:
+    """Return MatterAdapter instance."""
+    # NOTE: This assumes only one Matter connection/fabric can exist.
+    # Shall we support connecting to multiple servers in the client or by config entries?
+    # In case of the config entry we need to fix this.
+    matter_entry_data: MatterEntryData = next(iter(hass.data[DOMAIN].values()))
+    return matter_entry_data.adapter

homeassistant/components/matter/light.py

@@ -3,7 +3,7 @@ from __future__ import annotations
 
 from dataclasses import dataclass
 from functools import partial
-from typing import TYPE_CHECKING, Any
+from typing import Any
 
 from chip.clusters import Objects as clusters
 from matter_server.common.models import device_types
@@ -19,13 +19,10 @@ from homeassistant.const import Platform
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from .const import DOMAIN
 from .entity import MatterEntity, MatterEntityDescriptionBaseClass
+from .helpers import get_matter
 from .util import renormalize
 
-if TYPE_CHECKING:
-    from .adapter import MatterAdapter
-
 
 async def async_setup_entry(
     hass: HomeAssistant,
@@ -33,7 +30,7 @@ async def async_setup_entry(
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up Matter Light from Config Entry."""
-    matter: MatterAdapter = hass.data[DOMAIN][config_entry.entry_id]
+    matter = get_matter(hass)
     matter.register_platform_handler(Platform.LIGHT, async_add_entities)
 
 

homeassistant/components/matter/sensor.py

@@ -4,7 +4,7 @@ from __future__ import annotations
 from collections.abc import Callable
 from dataclasses import dataclass
 from functools import partial
-from typing import TYPE_CHECKING, Any
+from typing import Any
 
 from chip.clusters import Objects as clusters
 from chip.clusters.Types import Nullable, NullValue
@@ -29,11 +29,8 @@ from homeassistant.const import (
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from .const import DOMAIN
 from .entity import MatterEntity, MatterEntityDescriptionBaseClass
-
+from .helpers import get_matter
-if TYPE_CHECKING:
-    from .adapter import MatterAdapter
 
 
 async def async_setup_entry(
@@ -42,7 +39,7 @@ async def async_setup_entry(
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up Matter sensors from Config Entry."""
-    matter: MatterAdapter = hass.data[DOMAIN][config_entry.entry_id]
+    matter = get_matter(hass)
     matter.register_platform_handler(Platform.SENSOR, async_add_entities)
 
 

homeassistant/components/matter/switch.py

@@ -3,7 +3,7 @@ from __future__ import annotations
 
 from dataclasses import dataclass
 from functools import partial
-from typing import TYPE_CHECKING, Any
+from typing import Any
 
 from chip.clusters import Objects as clusters
 from matter_server.common.models import device_types
@@ -18,11 +18,8 @@ from homeassistant.const import Platform
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from .const import DOMAIN
 from .entity import MatterEntity, MatterEntityDescriptionBaseClass
-
+from .helpers import get_matter
-if TYPE_CHECKING:
-    from .adapter import MatterAdapter
 
 
 async def async_setup_entry(
@@ -31,7 +28,7 @@ async def async_setup_entry(
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up Matter switches from Config Entry."""
-    matter: MatterAdapter = hass.data[DOMAIN][config_entry.entry_id]
+    matter = get_matter(hass)
     matter.register_platform_handler(Platform.SWITCH, async_add_entities)
 
 

tests/components/matter/conftest.py

@@ -33,6 +33,9 @@ async def matter_client_fixture() -> AsyncGenerator[MagicMock, None]:
             """Mock listen."""
             if init_ready is not None:
                 init_ready.set()
+            listen_block = asyncio.Event()
+            await listen_block.wait()
+            assert False, "Listen was not cancelled!"
 
         client.connect = AsyncMock(side_effect=connect)
         client.start_listening = AsyncMock(side_effect=listen)

tests/components/matter/test_init.py

@@ -2,20 +2,211 @@
 from __future__ import annotations
 
 import asyncio
-from unittest.mock import AsyncMock, MagicMock, call
+from collections.abc import Generator
+from unittest.mock import AsyncMock, MagicMock, call, patch
 
-from matter_server.client.exceptions import InvalidServerVersion
+from matter_server.client.exceptions import CannotConnect, InvalidServerVersion
+from matter_server.common.helpers.util import dataclass_from_dict
+from matter_server.common.models.error import MatterError
+from matter_server.common.models.node import MatterNode
 import pytest
 
 from homeassistant.components.hassio import HassioAPIError
 from homeassistant.components.matter.const import DOMAIN
 from homeassistant.config_entries import ConfigEntryDisabler, ConfigEntryState
+from homeassistant.const import STATE_UNAVAILABLE
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers import issue_registry as ir
 
+from .common import load_and_parse_node_fixture
+
 from tests.common import MockConfigEntry
 
 
+@pytest.fixture(name="connect_timeout")
+def connect_timeout_fixture() -> Generator[int, None, None]:
+    """Mock the connect timeout."""
+    with patch("homeassistant.components.matter.CONNECT_TIMEOUT", new=0) as timeout:
+        yield timeout
+
+
+@pytest.fixture(name="listen_ready_timeout")
+def listen_ready_timeout_fixture() -> Generator[int, None, None]:
+    """Mock the listen ready timeout."""
+    with patch(
+        "homeassistant.components.matter.LISTEN_READY_TIMEOUT", new=0
+    ) as timeout:
+        yield timeout
+
+
+async def test_entry_setup_unload(
+    hass: HomeAssistant,
+    matter_client: MagicMock,
+) -> None:
+    """Test the integration set up and unload."""
+    node_data = load_and_parse_node_fixture("onoff-light")
+    node = dataclass_from_dict(
+        MatterNode,
+        node_data,
+    )
+    matter_client.get_nodes.return_value = [node]
+    matter_client.get_node.return_value = node
+    entry = MockConfigEntry(domain="matter", data={"url": "ws://localhost:5580/ws"})
+    entry.add_to_hass(hass)
+
+    await hass.config_entries.async_setup(entry.entry_id)
+    await hass.async_block_till_done()
+
+    assert matter_client.connect.call_count == 1
+    assert entry.state == ConfigEntryState.LOADED
+    entity_state = hass.states.get("light.mock_onoff_light")
+    assert entity_state
+    assert entity_state.state != STATE_UNAVAILABLE
+
+    await hass.config_entries.async_unload(entry.entry_id)
+
+    assert matter_client.disconnect.call_count == 1
+    assert entry.state == ConfigEntryState.NOT_LOADED
+    entity_state = hass.states.get("light.mock_onoff_light")
+    assert entity_state
+    assert entity_state.state == STATE_UNAVAILABLE
+
+
+async def test_home_assistant_stop(
+    hass: HomeAssistant,
+    matter_client: MagicMock,
+    integration: MockConfigEntry,
+) -> None:
+    """Test clean up on home assistant stop."""
+    await hass.async_stop()
+
+    assert matter_client.disconnect.call_count == 1
+
+
+@pytest.mark.parametrize("error", [CannotConnect("Boom"), Exception("Boom")])
+async def test_connect_failed(
+    hass: HomeAssistant,
+    matter_client: MagicMock,
+    error: Exception,
+) -> None:
+    """Test failure during client connection."""
+    entry = MockConfigEntry(domain=DOMAIN, data={"url": "ws://localhost:5580/ws"})
+    entry.add_to_hass(hass)
+    matter_client.connect.side_effect = error
+
+    await hass.config_entries.async_setup(entry.entry_id)
+    await hass.async_block_till_done()
+
+    assert entry.state is ConfigEntryState.SETUP_RETRY
+
+
+async def test_connect_timeout(
+    hass: HomeAssistant,
+    matter_client: MagicMock,
+    connect_timeout: int,
+) -> None:
+    """Test timeout during client connection."""
+    entry = MockConfigEntry(domain=DOMAIN, data={"url": "ws://localhost:5580/ws"})
+    entry.add_to_hass(hass)
+
+    await hass.config_entries.async_setup(entry.entry_id)
+    await hass.async_block_till_done()
+
+    assert entry.state is ConfigEntryState.SETUP_RETRY
+
+
+@pytest.mark.parametrize("error", [MatterError("Boom"), Exception("Boom")])
+async def test_listen_failure_timeout(
+    hass: HomeAssistant,
+    listen_ready_timeout: int,
+    matter_client: MagicMock,
+    error: Exception,
+) -> None:
+    """Test client listen errors during the first timeout phase."""
+
+    async def start_listening(listen_ready: asyncio.Event) -> None:
+        """Mock the client start_listening method."""
+        # Set the connect side effect to stop an endless loop on reload.
+        matter_client.connect.side_effect = MatterError("Boom")
+        raise error
+
+    matter_client.start_listening.side_effect = start_listening
+    entry = MockConfigEntry(domain=DOMAIN, data={"url": "ws://localhost:5580/ws"})
+    entry.add_to_hass(hass)
+
+    await hass.config_entries.async_setup(entry.entry_id)
+    await hass.async_block_till_done()
+
+    assert entry.state is ConfigEntryState.SETUP_RETRY
+
+
+@pytest.mark.parametrize("error", [MatterError("Boom"), Exception("Boom")])
+async def test_listen_failure_config_entry_not_loaded(
+    hass: HomeAssistant,
+    matter_client: MagicMock,
+    error: Exception,
+) -> None:
+    """Test client listen errors during the final phase before config entry loaded."""
+    listen_block = asyncio.Event()
+
+    async def start_listening(listen_ready: asyncio.Event) -> None:
+        """Mock the client start_listening method."""
+        listen_ready.set()
+        await listen_block.wait()
+        # Set the connect side effect to stop an endless loop on reload.
+        matter_client.connect.side_effect = MatterError("Boom")
+        raise error
+
+    async def get_nodes() -> list[MagicMock]:
+        """Mock the client get_nodes method."""
+        listen_block.set()
+        return []
+
+    matter_client.start_listening.side_effect = start_listening
+    matter_client.get_nodes.side_effect = get_nodes
+    entry = MockConfigEntry(domain=DOMAIN, data={"url": "ws://localhost:5580/ws"})
+    entry.add_to_hass(hass)
+
+    await hass.config_entries.async_setup(entry.entry_id)
+    await hass.async_block_till_done()
+
+    assert entry.state is ConfigEntryState.SETUP_RETRY
+    assert matter_client.disconnect.call_count == 1
+
+
+@pytest.mark.parametrize("error", [MatterError("Boom"), Exception("Boom")])
+async def test_listen_failure_config_entry_loaded(
+    hass: HomeAssistant,
+    matter_client: MagicMock,
+    error: Exception,
+) -> None:
+    """Test client listen errors after config entry is loaded."""
+    listen_block = asyncio.Event()
+
+    async def start_listening(listen_ready: asyncio.Event) -> None:
+        """Mock the client start_listening method."""
+        listen_ready.set()
+        await listen_block.wait()
+        # Set the connect side effect to stop an endless loop on reload.
+        matter_client.connect.side_effect = MatterError("Boom")
+        raise error
+
+    matter_client.start_listening.side_effect = start_listening
+    entry = MockConfigEntry(domain=DOMAIN, data={"url": "ws://localhost:5580/ws"})
+    entry.add_to_hass(hass)
+
+    await hass.config_entries.async_setup(entry.entry_id)
+    await hass.async_block_till_done()
+
+    assert entry.state == ConfigEntryState.LOADED
+
+    listen_block.set()
+    await hass.async_block_till_done()
+
+    assert entry.state == ConfigEntryState.SETUP_RETRY
+    assert matter_client.disconnect.call_count == 1
+
+
 async def test_raise_addon_task_in_progress(
     hass: HomeAssistant,
     addon_not_installed: AsyncMock,