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
--
GitLab