From 916c44b5b4d69eebba83241829b55d1815ee0b57 Mon Sep 17 00:00:00 2001
From: Nathan Spencer <natekspencer@gmail.com>
Date: Fri, 2 Sep 2022 14:18:10 -0600
Subject: [PATCH] Adjust litterrobot platform loading/unloading (#77682)

Co-authored-by: epenet <6771947+epenet@users.noreply.github.com>
---
 .../components/litterrobot/__init__.py        | 57 +++++++------------
 tests/components/litterrobot/conftest.py      |  4 +-
 tests/components/litterrobot/test_vacuum.py   |  9 ++-
 3 files changed, 30 insertions(+), 40 deletions(-)

diff --git a/homeassistant/components/litterrobot/__init__.py b/homeassistant/components/litterrobot/__init__.py
index d302989fc01..742e9dcb9c7 100644
--- a/homeassistant/components/litterrobot/__init__.py
+++ b/homeassistant/components/litterrobot/__init__.py
@@ -1,7 +1,7 @@
 """The Litter-Robot integration."""
 from __future__ import annotations
 
-from pylitterbot import FeederRobot, LitterRobot, LitterRobot3, LitterRobot4
+from pylitterbot import FeederRobot, LitterRobot, LitterRobot3, Robot
 
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import Platform
@@ -10,65 +10,48 @@ from homeassistant.core import HomeAssistant
 from .const import DOMAIN
 from .hub import LitterRobotHub
 
-PLATFORMS = [
-    Platform.BUTTON,
-    Platform.SELECT,
-    Platform.SENSOR,
-    Platform.SWITCH,
-    Platform.VACUUM,
-]
-
 PLATFORMS_BY_TYPE = {
-    LitterRobot: (
-        Platform.SELECT,
-        Platform.SENSOR,
-        Platform.SWITCH,
-        Platform.VACUUM,
-    ),
-    LitterRobot3: (
-        Platform.BUTTON,
-        Platform.SELECT,
-        Platform.SENSOR,
-        Platform.SWITCH,
-        Platform.VACUUM,
-    ),
-    LitterRobot4: (
-        Platform.SELECT,
-        Platform.SENSOR,
-        Platform.SWITCH,
-        Platform.VACUUM,
-    ),
-    FeederRobot: (
-        Platform.BUTTON,
+    Robot: (
         Platform.SELECT,
         Platform.SENSOR,
         Platform.SWITCH,
     ),
+    LitterRobot: (Platform.VACUUM,),
+    LitterRobot3: (Platform.BUTTON,),
+    FeederRobot: (Platform.BUTTON,),
 }
 
 
+def get_platforms_for_robots(robots: list[Robot]) -> set[Platform]:
+    """Get platforms for robots."""
+    return {
+        platform
+        for robot in robots
+        for robot_type, platforms in PLATFORMS_BY_TYPE.items()
+        if isinstance(robot, robot_type)
+        for platform in platforms
+    }
+
+
 async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Set up Litter-Robot from a config entry."""
     hass.data.setdefault(DOMAIN, {})
     hub = hass.data[DOMAIN][entry.entry_id] = LitterRobotHub(hass, entry.data)
     await hub.login(load_robots=True)
 
-    platforms: set[str] = set()
-    for robot in hub.account.robots:
-        platforms.update(PLATFORMS_BY_TYPE[type(robot)])
-    if platforms:
+    if platforms := get_platforms_for_robots(hub.account.robots):
         await hass.config_entries.async_forward_entry_setups(entry, platforms)
-
     return True
 
 
 async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Unload a config entry."""
-    unload_ok = await hass.config_entries.async_unload_platforms(entry, PLATFORMS)
-
     hub: LitterRobotHub = hass.data[DOMAIN][entry.entry_id]
     await hub.account.disconnect()
 
+    platforms = get_platforms_for_robots(hub.account.robots)
+    unload_ok = await hass.config_entries.async_unload_platforms(entry, platforms)
+
     if unload_ok:
         hass.data[DOMAIN].pop(entry.entry_id)
 
diff --git a/tests/components/litterrobot/conftest.py b/tests/components/litterrobot/conftest.py
index d5d29e12988..e5d5e730b61 100644
--- a/tests/components/litterrobot/conftest.py
+++ b/tests/components/litterrobot/conftest.py
@@ -99,8 +99,8 @@ async def setup_integration(
     with patch(
         "homeassistant.components.litterrobot.hub.Account", return_value=mock_account
     ), patch(
-        "homeassistant.components.litterrobot.PLATFORMS",
-        [platform_domain] if platform_domain else [],
+        "homeassistant.components.litterrobot.PLATFORMS_BY_TYPE",
+        {Robot: (platform_domain,)} if platform_domain else {},
     ):
         await hass.config_entries.async_setup(entry.entry_id)
         await hass.async_block_till_done()
diff --git a/tests/components/litterrobot/test_vacuum.py b/tests/components/litterrobot/test_vacuum.py
index eb9a4c8c60b..08aa8b2399b 100644
--- a/tests/components/litterrobot/test_vacuum.py
+++ b/tests/components/litterrobot/test_vacuum.py
@@ -52,6 +52,7 @@ async def test_vacuum(hass: HomeAssistant, mock_account: MagicMock) -> None:
     assert ent_reg_entry.unique_id == VACUUM_UNIQUE_ID_OLD
 
     await setup_integration(hass, mock_account, PLATFORM_DOMAIN)
+    assert len(ent_reg.entities) == 1
     assert hass.services.has_service(DOMAIN, SERVICE_SET_SLEEP_MODE)
 
     vacuum = hass.states.get(VACUUM_ENTITY_ID)
@@ -78,10 +79,16 @@ async def test_no_robots(
     hass: HomeAssistant, mock_account_with_no_robots: MagicMock
 ) -> None:
     """Tests the vacuum entity was set up."""
-    await setup_integration(hass, mock_account_with_no_robots, PLATFORM_DOMAIN)
+    entry = await setup_integration(hass, mock_account_with_no_robots, PLATFORM_DOMAIN)
 
     assert not hass.services.has_service(DOMAIN, SERVICE_SET_SLEEP_MODE)
 
+    ent_reg = er.async_get(hass)
+    assert len(ent_reg.entities) == 0
+
+    assert await hass.config_entries.async_unload(entry.entry_id)
+    await hass.async_block_till_done()
+
 
 async def test_vacuum_with_error(
     hass: HomeAssistant, mock_account_with_error: MagicMock