diff --git a/homeassistant/components/litterrobot/__init__.py b/homeassistant/components/litterrobot/__init__.py index 03435b300191..510da8435057 100644 --- a/homeassistant/components/litterrobot/__init__.py +++ b/homeassistant/components/litterrobot/__init__.py @@ -6,8 +6,13 @@ import logging from pylitterbot import Account from pylitterbot.exceptions import LitterRobotException -from homeassistant.const import CONF_PASSWORD, CONF_USERNAME, Platform -from homeassistant.core import HomeAssistant +from homeassistant.const import ( + CONF_PASSWORD, + CONF_USERNAME, + EVENT_HOMEASSISTANT_STOP, + Platform, +) +from homeassistant.core import Event, HomeAssistant from homeassistant.helpers import config_validation as cv from homeassistant.helpers.aiohttp_client import async_get_clientsession from homeassistant.helpers.device_registry import AnyDeviceEntry @@ -82,6 +87,16 @@ async def async_migrate_entry( async def async_setup_entry(hass: HomeAssistant, entry: LitterRobotConfigEntry) -> bool: """Set up Litter-Robot from a config entry.""" coordinator = LitterRobotDataUpdateCoordinator(hass, entry) + + # Entries are not unloaded at shutdown, and the first refresh already starts + # the account's WebSocket monitor. + async def _async_disconnect_account(event: Event) -> None: + await coordinator.account.disconnect() + + entry.async_on_unload( + hass.bus.async_listen_once(EVENT_HOMEASSISTANT_STOP, _async_disconnect_account) + ) + await coordinator.async_config_entry_first_refresh() entry.runtime_data = coordinator await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS) diff --git a/tests/components/litterrobot/test_init.py b/tests/components/litterrobot/test_init.py index 15125017acc4..0fb7a5ab9240 100644 --- a/tests/components/litterrobot/test_init.py +++ b/tests/components/litterrobot/test_init.py @@ -1,6 +1,7 @@ """Test Litter-Robot setup process.""" from datetime import timedelta +from typing import Any from unittest.mock import MagicMock, patch from freezegun.api import FrozenDateTimeFactory @@ -15,7 +16,11 @@ from homeassistant.components.vacuum import ( VacuumActivity, ) from homeassistant.config_entries import SOURCE_REAUTH, ConfigEntryState -from homeassistant.const import ATTR_ENTITY_ID, STATE_UNAVAILABLE +from homeassistant.const import ( + ATTR_ENTITY_ID, + EVENT_HOMEASSISTANT_STOP, + STATE_UNAVAILABLE, +) from homeassistant.core import HomeAssistant from homeassistant.helpers import device_registry as dr, entity_registry as er from homeassistant.setup import async_setup_component @@ -46,6 +51,34 @@ async def test_unload_entry(hass: HomeAssistant, mock_account: MagicMock) -> Non assert await hass.config_entries.async_unload(entry.entry_id) +async def test_shutdown_disconnects_account( + hass: HomeAssistant, mock_account: MagicMock +) -> None: + """Test the account is disconnected when Home Assistant stops.""" + await setup_integration(hass, mock_account, VACUUM_DOMAIN) + + hass.bus.async_fire(EVENT_HOMEASSISTANT_STOP) + await hass.async_block_till_done() + + mock_account.disconnect.assert_awaited_once() + + +async def test_shutdown_during_first_refresh_disconnects_account( + hass: HomeAssistant, mock_account: MagicMock +) -> None: + """Test a stop during the first refresh still disconnects the account.""" + + async def _stop_during_refresh(**kwargs: Any) -> None: + hass.bus.async_fire(EVENT_HOMEASSISTANT_STOP) + + mock_account.load_robots.side_effect = _stop_during_refresh + + await setup_integration(hass, mock_account, VACUUM_DOMAIN) + await hass.async_block_till_done() + + mock_account.disconnect.assert_awaited_once() + + @pytest.mark.parametrize( ("side_effect", "expected_state"), [