Disconnect the Litter-Robot account when Home Assistant stops (#180048)

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)

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"),
     [