Ensure Roborock disconnects mqtt on unload/stop (home-assistant#158144)

Co-authored-by: Copilot <[email protected]>

Author: allenporter

homeassistant/components/roborock/__init__.py

@@ -20,7 +20,7 @@
 from roborock.map.map_parser import MapParserConfig
 
 from homeassistant.const import CONF_USERNAME, EVENT_HOMEASSISTANT_STOP
-from homeassistant.core import HomeAssistant
+from homeassistant.core import Event, HomeAssistant
 from homeassistant.exceptions import ConfigEntryAuthFailed, ConfigEntryNotReady
 from homeassistant.helpers import device_registry as dr
 from homeassistant.helpers.aiohttp_client import async_get_clientsession
@@ -99,10 +99,17 @@ async def async_setup_entry(hass: HomeAssistant, entry: RoborockConfigEntry) ->
             translation_domain=DOMAIN,
             translation_key="home_data_fail",
         ) from err
+
+    async def shutdown_roborock(_: Event | None = None) -> None:
+        await asyncio.gather(device_manager.close(), cache.flush())
+
+    entry.async_on_unload(
+        hass.bus.async_listen_once(EVENT_HOMEASSISTANT_STOP, shutdown_roborock)
+    )
+    entry.async_on_unload(shutdown_roborock)
+
     devices = await device_manager.get_devices()
     _LOGGER.debug("Device manager found %d devices", len(devices))
-    for device in devices:
-        entry.async_on_unload(device.close)
 
     coordinators = await asyncio.gather(
         *build_setup_functions(hass, entry, devices, user_data),
@@ -124,25 +131,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: RoborockConfigEntry) ->
             translation_domain=DOMAIN,
             translation_key="no_coordinators",
         )
-    valid_coordinators = RoborockCoordinators(v1_coords, a01_coords)
-
-    async def on_stop(_: Any) -> None:
-        _LOGGER.debug("Shutting down roborock")
-        await asyncio.gather(
-            *(
-                coordinator.async_shutdown()
-                for coordinator in valid_coordinators.values()
-            ),
-            cache.flush(),
-        )
-
-    entry.async_on_unload(
-        hass.bus.async_listen_once(
-            EVENT_HOMEASSISTANT_STOP,
-            on_stop,
-        )
-    )
-    entry.runtime_data = valid_coordinators
+    entry.runtime_data = RoborockCoordinators(v1_coords, a01_coords)
 
     await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
 

tests/components/roborock/conftest.py

@@ -25,6 +25,7 @@
     ZeoState,
 )
 from roborock.devices.device import RoborockDevice
+from roborock.devices.device_manager import DeviceManager
 from roborock.devices.traits.v1 import PropertiesApi
 from roborock.devices.traits.v1.clean_summary import CleanSummaryTrait
 from roborock.devices.traits.v1.command import CommandTrait
@@ -134,18 +135,6 @@ async def close(self) -> None:
         """Close the device."""
 
 
-class FakeDeviceManager:
-    """A fake device manager that returns a list of devices."""
-
-    def __init__(self, devices: list[RoborockDevice]) -> None:
-        """Initialize the fake device manager."""
-        self._devices = devices
-
-    async def get_devices(self) -> list[RoborockDevice]:
-        """Return the list of devices."""
-        return self._devices
-
-
 def make_mock_trait(
     trait_spec: type[V1TraitMixin] | None = None,
     dataclass_template: RoborockBase | None = None,
@@ -348,16 +337,26 @@ def fake_vacuum_command_fixture(
     return command_trait
 
 
+@pytest.fixture(name="device_manager")
+def device_manager_fixture(
+    fake_devices: list[FakeDevice],
+) -> AsyncMock:
+    """Fixture to create a fake device manager."""
+    device_manager = AsyncMock(spec=DeviceManager)
+    device_manager.get_devices = AsyncMock(return_value=fake_devices)
+    return device_manager
+
+
 @pytest.fixture(name="fake_create_device_manager", autouse=True)
 def fake_create_device_manager_fixture(
-    fake_devices: list[FakeDevice],
-) -> Generator[Mock]:
+    device_manager: AsyncMock,
+) -> None:
     """Fixture to create a fake device manager."""
     with patch(
         "homeassistant.components.roborock.create_device_manager",
     ) as mock_create_device_manager:
-        mock_create_device_manager.return_value = FakeDeviceManager(fake_devices)
-        yield mock_create_device_manager
+        mock_create_device_manager.return_value = device_manager
+        yield
 
 
 @pytest.fixture(name="config_entry_data")

tests/components/roborock/test_init.py

@@ -2,7 +2,7 @@
 
 import pathlib
 from typing import Any
-from unittest.mock import patch
+from unittest.mock import AsyncMock, patch
 
 import pytest
 from roborock import (
@@ -26,14 +26,42 @@
 from tests.typing import ClientSessionGenerator
 
 
-async def test_unload_entry(hass: HomeAssistant, setup_entry: MockConfigEntry) -> None:
-    """Test unloading roboorck integration."""
+async def test_unload_entry(
+    hass: HomeAssistant,
+    setup_entry: MockConfigEntry,
+    device_manager: AsyncMock,
+) -> None:
+    """Test unloading roborock integration."""
     assert len(hass.config_entries.async_entries(DOMAIN)) == 1
     assert setup_entry.state is ConfigEntryState.LOADED
+
+    assert device_manager.get_devices.called
+    assert not device_manager.close.called
+
+    # Unload the config entry and verify that the device manager is closed
     assert await hass.config_entries.async_unload(setup_entry.entry_id)
     await hass.async_block_till_done()
     assert setup_entry.state is ConfigEntryState.NOT_LOADED
 
+    assert device_manager.close.called
+
+
+async def test_home_assistant_stop(
+    hass: HomeAssistant,
+    setup_entry: MockConfigEntry,
+    device_manager: AsyncMock,
+) -> None:
+    """Test shutting down Home Assistant."""
+    assert len(hass.config_entries.async_entries(DOMAIN)) == 1
+    assert setup_entry.state is ConfigEntryState.LOADED
+
+    assert not device_manager.close.called
+
+    # Perform Home Assistant stop and verify that device manager is closed
+    await hass.async_stop()
+
+    assert device_manager.close.called
+
 
 async def test_reauth_started(
     hass: HomeAssistant, mock_roborock_entry: MockConfigEntry