Adjust litterrobot platform loading/unloading (#77682)
Co-authored-by: epenet <6771947+epenet@users.noreply.github.com>
This commit is contained in:
parent
e4e29aa29d
commit
916c44b5b4
3 changed files with 30 additions and 40 deletions
|
@ -1,7 +1,7 @@
|
||||||
"""The Litter-Robot integration."""
|
"""The Litter-Robot integration."""
|
||||||
from __future__ import annotations
|
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.config_entries import ConfigEntry
|
||||||
from homeassistant.const import Platform
|
from homeassistant.const import Platform
|
||||||
|
@ -10,40 +10,26 @@ from homeassistant.core import HomeAssistant
|
||||||
from .const import DOMAIN
|
from .const import DOMAIN
|
||||||
from .hub import LitterRobotHub
|
from .hub import LitterRobotHub
|
||||||
|
|
||||||
PLATFORMS = [
|
|
||||||
Platform.BUTTON,
|
|
||||||
Platform.SELECT,
|
|
||||||
Platform.SENSOR,
|
|
||||||
Platform.SWITCH,
|
|
||||||
Platform.VACUUM,
|
|
||||||
]
|
|
||||||
|
|
||||||
PLATFORMS_BY_TYPE = {
|
PLATFORMS_BY_TYPE = {
|
||||||
LitterRobot: (
|
Robot: (
|
||||||
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.SELECT,
|
||||||
Platform.SENSOR,
|
Platform.SENSOR,
|
||||||
Platform.SWITCH,
|
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
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
||||||
|
@ -53,22 +39,19 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
|
||||||
hub = hass.data[DOMAIN][entry.entry_id] = LitterRobotHub(hass, entry.data)
|
hub = hass.data[DOMAIN][entry.entry_id] = LitterRobotHub(hass, entry.data)
|
||||||
await hub.login(load_robots=True)
|
await hub.login(load_robots=True)
|
||||||
|
|
||||||
platforms: set[str] = set()
|
if platforms := get_platforms_for_robots(hub.account.robots):
|
||||||
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)
|
await hass.config_entries.async_forward_entry_setups(entry, platforms)
|
||||||
|
|
||||||
return True
|
return True
|
||||||
|
|
||||||
|
|
||||||
async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
|
async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
|
||||||
"""Unload a config entry."""
|
"""Unload a config entry."""
|
||||||
unload_ok = await hass.config_entries.async_unload_platforms(entry, PLATFORMS)
|
|
||||||
|
|
||||||
hub: LitterRobotHub = hass.data[DOMAIN][entry.entry_id]
|
hub: LitterRobotHub = hass.data[DOMAIN][entry.entry_id]
|
||||||
await hub.account.disconnect()
|
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:
|
if unload_ok:
|
||||||
hass.data[DOMAIN].pop(entry.entry_id)
|
hass.data[DOMAIN].pop(entry.entry_id)
|
||||||
|
|
||||||
|
|
|
@ -99,8 +99,8 @@ async def setup_integration(
|
||||||
with patch(
|
with patch(
|
||||||
"homeassistant.components.litterrobot.hub.Account", return_value=mock_account
|
"homeassistant.components.litterrobot.hub.Account", return_value=mock_account
|
||||||
), patch(
|
), patch(
|
||||||
"homeassistant.components.litterrobot.PLATFORMS",
|
"homeassistant.components.litterrobot.PLATFORMS_BY_TYPE",
|
||||||
[platform_domain] if platform_domain else [],
|
{Robot: (platform_domain,)} if platform_domain else {},
|
||||||
):
|
):
|
||||||
await hass.config_entries.async_setup(entry.entry_id)
|
await hass.config_entries.async_setup(entry.entry_id)
|
||||||
await hass.async_block_till_done()
|
await hass.async_block_till_done()
|
||||||
|
|
|
@ -52,6 +52,7 @@ async def test_vacuum(hass: HomeAssistant, mock_account: MagicMock) -> None:
|
||||||
assert ent_reg_entry.unique_id == VACUUM_UNIQUE_ID_OLD
|
assert ent_reg_entry.unique_id == VACUUM_UNIQUE_ID_OLD
|
||||||
|
|
||||||
await setup_integration(hass, mock_account, PLATFORM_DOMAIN)
|
await setup_integration(hass, mock_account, PLATFORM_DOMAIN)
|
||||||
|
assert len(ent_reg.entities) == 1
|
||||||
assert hass.services.has_service(DOMAIN, SERVICE_SET_SLEEP_MODE)
|
assert hass.services.has_service(DOMAIN, SERVICE_SET_SLEEP_MODE)
|
||||||
|
|
||||||
vacuum = hass.states.get(VACUUM_ENTITY_ID)
|
vacuum = hass.states.get(VACUUM_ENTITY_ID)
|
||||||
|
@ -78,10 +79,16 @@ async def test_no_robots(
|
||||||
hass: HomeAssistant, mock_account_with_no_robots: MagicMock
|
hass: HomeAssistant, mock_account_with_no_robots: MagicMock
|
||||||
) -> None:
|
) -> None:
|
||||||
"""Tests the vacuum entity was set up."""
|
"""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)
|
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(
|
async def test_vacuum_with_error(
|
||||||
hass: HomeAssistant, mock_account_with_error: MagicMock
|
hass: HomeAssistant, mock_account_with_error: MagicMock
|
||||||
|
|
Loading…
Add table
Reference in a new issue