From 2ac15ab67d5419bda5f0a4f0cac2a3b2028afbbb Mon Sep 17 00:00:00 2001 From: Allen Porter Date: Mon, 8 Dec 2025 00:36:34 -0800 Subject: [PATCH] Ensure Roborock disconnects mqtt on unload/stop (#158144) Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- homeassistant/components/roborock/__init__.py | 33 ++++++------------ tests/components/roborock/conftest.py | 31 ++++++++--------- tests/components/roborock/test_init.py | 34 +++++++++++++++++-- 3 files changed, 57 insertions(+), 41 deletions(-) diff --git a/homeassistant/components/roborock/__init__.py b/homeassistant/components/roborock/__init__.py index ece5a62e48d3..5d67233f7bc7 100644 --- a/homeassistant/components/roborock/__init__.py +++ b/homeassistant/components/roborock/__init__.py @@ -20,7 +20,7 @@ from roborock.devices.device_manager import UserParams, create_device_manager 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) diff --git a/tests/components/roborock/conftest.py b/tests/components/roborock/conftest.py index ef568cad8e04..17637b3bc4d6 100644 --- a/tests/components/roborock/conftest.py +++ b/tests/components/roborock/conftest.py @@ -25,6 +25,7 @@ from roborock.data import ( 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 @@ class FakeDevice(RoborockDevice): """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") diff --git a/tests/components/roborock/test_init.py b/tests/components/roborock/test_init.py index 66932405105e..4773badbe14f 100644 --- a/tests/components/roborock/test_init.py +++ b/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.common import MockConfigEntry 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