1
0
mirror of https://github.com/home-assistant/core.git synced 2026-03-01 14:25:31 +00:00
Files
core/tests/components/vacuum/__init__.py
Artur Pragacz 7a41ce1fd8 Add clean_area action to vacuum (#149315)
Co-authored-by: Bram Kragten <mail@bramkragten.nl>
2026-02-18 14:13:08 +01:00

128 lines
3.8 KiB
Python

"""The tests for vacuum platforms."""
from typing import Any
from homeassistant.components.vacuum import (
Segment,
StateVacuumEntity,
VacuumActivity,
VacuumEntityFeature,
)
from homeassistant.config_entries import ConfigEntry
from homeassistant.const import Platform
from homeassistant.core import HomeAssistant
from tests.common import MockEntity
class MockVacuum(MockEntity, StateVacuumEntity):
"""Mock vacuum class."""
_attr_supported_features = (
VacuumEntityFeature.PAUSE
| VacuumEntityFeature.STOP
| VacuumEntityFeature.RETURN_HOME
| VacuumEntityFeature.FAN_SPEED
| VacuumEntityFeature.BATTERY
| VacuumEntityFeature.CLEAN_SPOT
| VacuumEntityFeature.MAP
| VacuumEntityFeature.STATE
| VacuumEntityFeature.START
)
_attr_battery_level = 99
_attr_fan_speed_list = ["slow", "fast"]
def __init__(self, **values: Any) -> None:
"""Initialize a mock vacuum entity."""
super().__init__(**values)
self._attr_activity = VacuumActivity.DOCKED
self._attr_fan_speed = "slow"
def stop(self, **kwargs: Any) -> None:
"""Stop cleaning."""
self._attr_activity = VacuumActivity.IDLE
def return_to_base(self, **kwargs: Any) -> None:
"""Return to base."""
self._attr_activity = VacuumActivity.RETURNING
def clean_spot(self, **kwargs: Any) -> None:
"""Clean a spot."""
self._attr_activity = VacuumActivity.CLEANING
def set_fan_speed(self, fan_speed: str, **kwargs: Any) -> None:
"""Set the fan speed."""
self._attr_fan_speed = fan_speed
def start(self) -> None:
"""Start cleaning."""
self._attr_activity = VacuumActivity.CLEANING
def pause(self) -> None:
"""Pause cleaning."""
self._attr_activity = VacuumActivity.PAUSED
async def help_async_setup_entry_init(
hass: HomeAssistant, config_entry: ConfigEntry
) -> bool:
"""Set up test config entry."""
await hass.config_entries.async_forward_entry_setups(
config_entry, [Platform.VACUUM]
)
return True
async def help_async_unload_entry(
hass: HomeAssistant, config_entry: ConfigEntry
) -> bool:
"""Unload test config emntry."""
return await hass.config_entries.async_unload_platforms(
config_entry, [Platform.VACUUM]
)
SEGMENTS = [
Segment(id="seg_1", name="Kitchen"),
Segment(id="seg_2", name="Living Room"),
Segment(id="seg_3", name="Bedroom"),
Segment(id="seg_4", name="Bedroom", group="Upstairs"),
Segment(id="seg_5", name="Bathroom", group="Upstairs"),
]
class MockVacuumWithCleanArea(MockEntity, StateVacuumEntity):
"""Mock vacuum with clean_area support."""
_attr_supported_features = (
VacuumEntityFeature.STATE
| VacuumEntityFeature.START
| VacuumEntityFeature.CLEAN_AREA
)
def __init__(
self,
segments: list[Segment] | None = None,
unique_id: str = "mock_vacuum_unique_id",
**values: Any,
) -> None:
"""Initialize a mock vacuum entity."""
super().__init__(**values)
self._attr_unique_id = unique_id
self._attr_activity = VacuumActivity.DOCKED
self.segments = segments if segments is not None else SEGMENTS
self.clean_segments_calls: list[tuple[list[str], dict[str, Any]]] = []
def start(self) -> None:
"""Start cleaning."""
self._attr_activity = VacuumActivity.CLEANING
async def async_get_segments(self) -> list[Segment]:
"""Get the segments that can be cleaned."""
return self.segments
async def async_clean_segments(self, segment_ids: list[str], **kwargs: Any) -> None:
"""Perform an area clean."""
self.clean_segments_calls.append((segment_ids, kwargs))
self._attr_activity = VacuumActivity.CLEANING