Let platforms decide entity creation in litterrobot (#136738)

homeassistant/components/litterrobot/__init__.py

@@ -2,8 +2,6 @@
 
 from __future__ import annotations
 
-from pylitterbot import FeederRobot, LitterRobot, LitterRobot3, LitterRobot4, Robot
-
 from homeassistant.const import Platform
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.device_registry import DeviceEntry
@@ -11,29 +9,16 @@ from homeassistant.helpers.device_registry import DeviceEntry
 from .const import DOMAIN
 from .coordinator import LitterRobotConfigEntry, LitterRobotDataUpdateCoordinator
 
-PLATFORMS_BY_TYPE = {
+PLATFORMS = [
-    Robot: (
     Platform.BINARY_SENSOR,
+    Platform.BUTTON,
     Platform.SELECT,
     Platform.SENSOR,
     Platform.SWITCH,
-    ),
+    Platform.TIME,
-    LitterRobot: (Platform.VACUUM,),
+    Platform.UPDATE,
-    LitterRobot3: (Platform.BUTTON, Platform.TIME),
+    Platform.VACUUM,
-    LitterRobot4: (Platform.UPDATE,),
+]
-    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: LitterRobotConfigEntry) -> bool:
@@ -41,9 +26,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: LitterRobotConfigEntry)
     coordinator = LitterRobotDataUpdateCoordinator(hass, entry)
     await coordinator.async_config_entry_first_refresh()
     entry.runtime_data = coordinator
-
+    await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
-    if platforms := get_platforms_for_robots(coordinator.account.robots):
-        await hass.config_entries.async_forward_entry_setups(entry, platforms)
     return True
 
 
@@ -52,9 +35,7 @@ async def async_unload_entry(
 ) -> bool:
     """Unload a config entry."""
     await entry.runtime_data.account.disconnect()
-
+    return await hass.config_entries.async_unload_platforms(entry, PLATFORMS)
-    platforms = get_platforms_for_robots(entry.runtime_data.account.robots)
-    return await hass.config_entries.async_unload_platforms(entry, platforms)
 
 
 async def async_remove_config_entry_device(

tests/components/litterrobot/conftest.py

@@ -123,15 +123,9 @@ async def setup_integration(
     )
     entry.add_to_hass(hass)
 
-    with (
+    with patch(
-        patch(
         "homeassistant.components.litterrobot.coordinator.Account",
         return_value=mock_account,
-        ),
-        patch(
-            "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()

tests/components/litterrobot/test_vacuum.py

@@ -33,7 +33,6 @@ async def test_vacuum(
     hass: HomeAssistant, entity_registry: er.EntityRegistry, mock_account: MagicMock
 ) -> None:
     """Tests the vacuum entity was set up."""
-
     entity_registry.async_get_or_create(
         VACUUM_DOMAIN,
         DOMAIN,
@@ -44,7 +43,6 @@ async def test_vacuum(
     assert ent_reg_entry.unique_id == VACUUM_UNIQUE_ID
 
     await setup_integration(hass, mock_account, VACUUM_DOMAIN)
-    assert len(entity_registry.entities) == 1
     assert hass.services.has_service(DOMAIN, SERVICE_SET_SLEEP_MODE)
 
     vacuum = hass.states.get(VACUUM_ENTITY_ID)
@@ -63,8 +61,6 @@ async def test_no_robots(
     """Tests the vacuum entity was set up."""
     entry = await setup_integration(hass, mock_account_with_no_robots, VACUUM_DOMAIN)
 
-    assert not hass.services.has_service(DOMAIN, SERVICE_SET_SLEEP_MODE)
-
     assert len(entity_registry.entities) == 0
 
     assert await hass.config_entries.async_unload(entry.entry_id)