diff --git a/homeassistant/components/litterrobot/__init__.py b/homeassistant/components/litterrobot/__init__.py index 5aa186d0171..d302989fc01 100644 --- a/homeassistant/components/litterrobot/__init__.py +++ b/homeassistant/components/litterrobot/__init__.py @@ -1,6 +1,8 @@ """The Litter-Robot integration.""" from __future__ import annotations +from pylitterbot import FeederRobot, LitterRobot, LitterRobot3, LitterRobot4 + from homeassistant.config_entries import ConfigEntry from homeassistant.const import Platform from homeassistant.core import HomeAssistant @@ -16,6 +18,34 @@ PLATFORMS = [ 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, + Platform.SELECT, + Platform.SENSOR, + Platform.SWITCH, + ), +} + async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: """Set up Litter-Robot from a config entry.""" @@ -23,8 +53,11 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: hub = hass.data[DOMAIN][entry.entry_id] = LitterRobotHub(hass, entry.data) await hub.login(load_robots=True) - if any(hub.litter_robots()): - await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS) + platforms: set[str] = set() + for robot in hub.account.robots: + platforms.update(PLATFORMS_BY_TYPE[type(robot)]) + if platforms: + await hass.config_entries.async_forward_entry_setups(entry, platforms) return True