Add clean_area action to vacuum (#149315)

Co-authored-by: Bram Kragten <mail@bramkragten.nl>

tests/components/demo/test_vacuum.py

@@ -70,7 +70,7 @@ async def setup_demo_vacuum(hass: HomeAssistant, vacuum_only: None):
 async def test_supported_features(hass: HomeAssistant) -> None:
     """Test vacuum supported features."""
     state = hass.states.get(ENTITY_VACUUM_COMPLETE)
-    assert state.attributes.get(ATTR_SUPPORTED_FEATURES) == 16316
+    assert state.attributes.get(ATTR_SUPPORTED_FEATURES) == 32700
     assert state.attributes.get(ATTR_FAN_SPEED) == "medium"
     assert state.attributes.get(ATTR_FAN_SPEED_LIST) == FAN_SPEEDS
     assert state.state == VacuumActivity.DOCKED

tests/components/vacuum/__init__.py

@@ -3,6 +3,7 @@
 from typing import Any
 
 from homeassistant.components.vacuum import (
+    Segment,
     StateVacuumEntity,
     VacuumActivity,
     VacuumEntityFeature,
@@ -79,3 +80,48 @@ async def help_async_unload_entry(
     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

tests/components/vacuum/test_init.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+from dataclasses import asdict
 import logging
 from typing import Any
 
@@ -9,6 +10,7 @@ import pytest
 
 from homeassistant.components.vacuum import (
     DOMAIN,
+    SERVICE_CLEAN_AREA,
     SERVICE_CLEAN_SPOT,
     SERVICE_LOCATE,
     SERVICE_PAUSE,
@@ -22,12 +24,19 @@ from homeassistant.components.vacuum import (
     VacuumEntityFeature,
 )
 from homeassistant.core import HomeAssistant
+from homeassistant.helpers import entity_registry as er, issue_registry as ir
 
-from . import MockVacuum, help_async_setup_entry_init, help_async_unload_entry
+from . import (
+    MockVacuum,
+    MockVacuumWithCleanArea,
+    help_async_setup_entry_init,
+    help_async_unload_entry,
+)
 from .common import async_start
 
 from tests.common import (
     MockConfigEntry,
+    MockEntity,
     MockModule,
     mock_integration,
     setup_test_component_platform,
@@ -206,6 +215,252 @@ async def test_send_command(hass: HomeAssistant, config_flow_fixture: None) -> N
     assert "test" in strings
 
 
+@pytest.mark.usefixtures("config_flow_fixture")
+@pytest.mark.parametrize(
+    ("area_mapping", "targeted_areas", "targeted_segments"),
+    [
+        (
+            {"area_1": ["seg_1"], "area_2": ["seg_2", "seg_3"]},
+            ["area_1", "area_2"],
+            ["seg_1", "seg_2", "seg_3"],
+        ),
+        (
+            {"area_1": ["seg_1", "seg_2"], "area_2": ["seg_2", "seg_3"]},
+            ["area_1", "area_2"],
+            ["seg_1", "seg_2", "seg_3"],
+        ),
+    ],
+)
+async def test_clean_area_service(
+    hass: HomeAssistant,
+    entity_registry: er.EntityRegistry,
+    area_mapping: dict[str, list[str]],
+    targeted_areas: list[str],
+    targeted_segments: list[str],
+) -> None:
+    """Test clean_area service calls async_clean_segments with correct segments."""
+    mock_vacuum = MockVacuumWithCleanArea(name="Testing", entity_id="vacuum.testing")
+
+    config_entry = MockConfigEntry(domain="test")
+    config_entry.add_to_hass(hass)
+
+    mock_integration(
+        hass,
+        MockModule(
+            "test",
+            async_setup_entry=help_async_setup_entry_init,
+            async_unload_entry=help_async_unload_entry,
+        ),
+    )
+    setup_test_component_platform(hass, DOMAIN, [mock_vacuum], from_config_entry=True)
+    assert await hass.config_entries.async_setup(config_entry.entry_id)
+    await hass.async_block_till_done()
+
+    entity_registry.async_update_entity_options(
+        mock_vacuum.entity_id,
+        DOMAIN,
+        {
+            "area_mapping": area_mapping,
+            "last_seen_segments": [asdict(segment) for segment in mock_vacuum.segments],
+        },
+    )
+
+    await hass.services.async_call(
+        DOMAIN,
+        SERVICE_CLEAN_AREA,
+        {"entity_id": mock_vacuum.entity_id, "cleaning_area_id": targeted_areas},
+        blocking=True,
+    )
+
+    assert len(mock_vacuum.clean_segments_calls) == 1
+    assert mock_vacuum.clean_segments_calls[0][0] == targeted_segments
+
+
+@pytest.mark.usefixtures("config_flow_fixture")
+@pytest.mark.parametrize(
+    ("area_mapping", "targeted_areas"),
+    [
+        ({}, ["area_1"]),
+        ({"area_1": ["seg_1"]}, ["area_2"]),
+    ],
+)
+async def test_clean_area_no_segments(
+    hass: HomeAssistant,
+    entity_registry: er.EntityRegistry,
+    area_mapping: dict[str, list[str]],
+    targeted_areas: list[str],
+) -> None:
+    """Test clean_area does nothing when no segments to clean."""
+    mock_vacuum = MockVacuumWithCleanArea(name="Testing", entity_id="vacuum.testing")
+
+    config_entry = MockConfigEntry(domain="test")
+    config_entry.add_to_hass(hass)
+
+    mock_integration(
+        hass,
+        MockModule(
+            "test",
+            async_setup_entry=help_async_setup_entry_init,
+            async_unload_entry=help_async_unload_entry,
+        ),
+    )
+    setup_test_component_platform(hass, DOMAIN, [mock_vacuum], from_config_entry=True)
+    assert await hass.config_entries.async_setup(config_entry.entry_id)
+    await hass.async_block_till_done()
+
+    await hass.services.async_call(
+        DOMAIN,
+        SERVICE_CLEAN_AREA,
+        {"entity_id": mock_vacuum.entity_id, "cleaning_area_id": targeted_areas},
+        blocking=True,
+    )
+
+    entity_registry.async_update_entity_options(
+        mock_vacuum.entity_id,
+        DOMAIN,
+        {
+            "area_mapping": area_mapping,
+            "last_seen_segments": [asdict(segment) for segment in mock_vacuum.segments],
+        },
+    )
+
+    await hass.services.async_call(
+        DOMAIN,
+        SERVICE_CLEAN_AREA,
+        {"entity_id": mock_vacuum.entity_id, "cleaning_area_id": targeted_areas},
+        blocking=True,
+    )
+
+    assert len(mock_vacuum.clean_segments_calls) == 0
+
+
+@pytest.mark.usefixtures("config_flow_fixture")
+async def test_clean_area_methods_not_implemented(hass: HomeAssistant) -> None:
+    """Test async_get_segments and async_clean_segments raise NotImplementedError."""
+
+    class MockVacuumNoImpl(MockEntity, StateVacuumEntity):
+        """Mock vacuum without implementations."""
+
+        _attr_supported_features = (
+            VacuumEntityFeature.STATE | VacuumEntityFeature.CLEAN_AREA
+        )
+        _attr_activity = VacuumActivity.DOCKED
+
+    mock_vacuum = MockVacuumNoImpl(name="Testing", entity_id="vacuum.testing")
+
+    config_entry = MockConfigEntry(domain="test")
+    config_entry.add_to_hass(hass)
+
+    mock_integration(
+        hass,
+        MockModule(
+            "test",
+            async_setup_entry=help_async_setup_entry_init,
+            async_unload_entry=help_async_unload_entry,
+        ),
+    )
+    setup_test_component_platform(hass, DOMAIN, [mock_vacuum], from_config_entry=True)
+    assert await hass.config_entries.async_setup(config_entry.entry_id)
+    await hass.async_block_till_done()
+
+    with pytest.raises(NotImplementedError):
+        await mock_vacuum.async_get_segments()
+
+    with pytest.raises(NotImplementedError):
+        await mock_vacuum.async_clean_segments(["seg_1"])
+
+
+async def test_clean_area_no_registry_entry() -> None:
+    """Test error handling when registry entry is not set."""
+    mock_vacuum = MockVacuumWithCleanArea(name="Testing", entity_id="vacuum.testing")
+
+    with pytest.raises(
+        RuntimeError,
+        match="Cannot access last_seen_segments, registry entry is not set",
+    ):
+        mock_vacuum.last_seen_segments  # noqa: B018
+
+    with pytest.raises(
+        RuntimeError,
+        match="Cannot perform area clean, registry entry is not set",
+    ):
+        await mock_vacuum.async_internal_clean_area(["area_1"])
+
+    with pytest.raises(
+        RuntimeError,
+        match="Cannot create segments issue, registry entry is not set",
+    ):
+        mock_vacuum.async_create_segments_issue()
+
+
+@pytest.mark.usefixtures("config_flow_fixture")
+async def test_last_seen_segments(
+    hass: HomeAssistant, entity_registry: er.EntityRegistry
+) -> None:
+    """Test last_seen_segments property."""
+    mock_vacuum = MockVacuumWithCleanArea(name="Testing", entity_id="vacuum.testing")
+
+    config_entry = MockConfigEntry(domain="test")
+    config_entry.add_to_hass(hass)
+
+    mock_integration(
+        hass,
+        MockModule(
+            "test",
+            async_setup_entry=help_async_setup_entry_init,
+            async_unload_entry=help_async_unload_entry,
+        ),
+    )
+    setup_test_component_platform(hass, DOMAIN, [mock_vacuum], from_config_entry=True)
+    assert await hass.config_entries.async_setup(config_entry.entry_id)
+    await hass.async_block_till_done()
+
+    assert mock_vacuum.last_seen_segments is None
+
+    entity_registry.async_update_entity_options(
+        mock_vacuum.entity_id,
+        DOMAIN,
+        {
+            "area_mapping": {},
+            "last_seen_segments": [asdict(segment) for segment in mock_vacuum.segments],
+        },
+    )
+
+    assert mock_vacuum.last_seen_segments == mock_vacuum.segments
+
+
+@pytest.mark.usefixtures("config_flow_fixture")
+async def test_last_seen_segments_and_issue_creation(
+    hass: HomeAssistant, entity_registry: er.EntityRegistry
+) -> None:
+    """Test last_seen_segments property and segments issue creation."""
+    mock_vacuum = MockVacuumWithCleanArea(name="Testing", entity_id="vacuum.testing")
+
+    config_entry = MockConfigEntry(domain="test")
+    config_entry.add_to_hass(hass)
+
+    mock_integration(
+        hass,
+        MockModule(
+            "test",
+            async_setup_entry=help_async_setup_entry_init,
+            async_unload_entry=help_async_unload_entry,
+        ),
+    )
+    setup_test_component_platform(hass, DOMAIN, [mock_vacuum], from_config_entry=True)
+    assert await hass.config_entries.async_setup(config_entry.entry_id)
+    await hass.async_block_till_done()
+
+    entity_entry = entity_registry.async_get(mock_vacuum.entity_id)
+    mock_vacuum.async_create_segments_issue()
+
+    issue_id = f"segments_changed_{entity_entry.id}"
+    issue = ir.async_get(hass).async_get_issue(DOMAIN, issue_id)
+    assert issue is not None
+    assert issue.severity == ir.IssueSeverity.WARNING
+    assert issue.translation_key == "segments_changed"
+
+
 @pytest.mark.parametrize(("is_built_in", "log_warnings"), [(True, 0), (False, 3)])
 async def test_vacuum_log_deprecated_battery_using_properties(
     hass: HomeAssistant,

tests/components/vacuum/test_websocket.py

@@ -0,0 +1,125 @@
+"""Tests for vacuum websocket API."""
+
+from __future__ import annotations
+
+from dataclasses import asdict
+
+import pytest
+
+from homeassistant.components.vacuum import (
+    DOMAIN,
+    Segment,
+    StateVacuumEntity,
+    VacuumActivity,
+    VacuumEntityFeature,
+)
+from homeassistant.components.websocket_api import ERR_NOT_FOUND, ERR_NOT_SUPPORTED
+from homeassistant.core import HomeAssistant
+from homeassistant.setup import async_setup_component
+
+from . import (
+    MockVacuumWithCleanArea,
+    help_async_setup_entry_init,
+    help_async_unload_entry,
+)
+
+from tests.common import (
+    MockConfigEntry,
+    MockEntity,
+    MockModule,
+    mock_integration,
+    setup_test_component_platform,
+)
+from tests.typing import WebSocketGenerator
+
+
+@pytest.mark.usefixtures("config_flow_fixture")
+async def test_get_segments(
+    hass: HomeAssistant,
+    hass_ws_client: WebSocketGenerator,
+) -> None:
+    """Test vacuum/get_segments websocket command."""
+    segments = [
+        Segment(id="seg_1", name="Kitchen"),
+        Segment(id="seg_2", name="Living Room"),
+        Segment(id="seg_3", name="Bedroom", group="Upstairs"),
+    ]
+    entity = MockVacuumWithCleanArea(
+        name="Testing",
+        entity_id="vacuum.testing",
+        segments=segments,
+    )
+    config_entry = MockConfigEntry(domain="test")
+    config_entry.add_to_hass(hass)
+
+    mock_integration(
+        hass,
+        MockModule(
+            "test",
+            async_setup_entry=help_async_setup_entry_init,
+            async_unload_entry=help_async_unload_entry,
+        ),
+    )
+    setup_test_component_platform(hass, DOMAIN, [entity], from_config_entry=True)
+    assert await hass.config_entries.async_setup(config_entry.entry_id)
+    await hass.async_block_till_done()
+
+    client = await hass_ws_client(hass)
+    await client.send_json_auto_id(
+        {"type": "vacuum/get_segments", "entity_id": entity.entity_id}
+    )
+    msg = await client.receive_json()
+    assert msg["success"]
+    assert msg["result"] == {"segments": [asdict(seg) for seg in segments]}
+
+
+async def test_get_segments_entity_not_found(
+    hass: HomeAssistant,
+    hass_ws_client: WebSocketGenerator,
+) -> None:
+    """Test vacuum/get_segments with unknown entity."""
+    assert await async_setup_component(hass, DOMAIN, {})
+
+    client = await hass_ws_client(hass)
+    await client.send_json_auto_id(
+        {"type": "vacuum/get_segments", "entity_id": "vacuum.unknown"}
+    )
+    msg = await client.receive_json()
+    assert not msg["success"]
+    assert msg["error"]["code"] == ERR_NOT_FOUND
+
+
+@pytest.mark.usefixtures("config_flow_fixture")
+async def test_get_segments_not_supported(
+    hass: HomeAssistant,
+    hass_ws_client: WebSocketGenerator,
+) -> None:
+    """Test vacuum/get_segments with entity not supporting CLEAN_AREA."""
+
+    class MockVacuumNoCleanArea(MockEntity, StateVacuumEntity):
+        _attr_supported_features = VacuumEntityFeature.STATE | VacuumEntityFeature.START
+        _attr_activity = VacuumActivity.DOCKED
+
+    entity = MockVacuumNoCleanArea(name="Testing", entity_id="vacuum.testing")
+    config_entry = MockConfigEntry(domain="test")
+    config_entry.add_to_hass(hass)
+
+    mock_integration(
+        hass,
+        MockModule(
+            "test",
+            async_setup_entry=help_async_setup_entry_init,
+            async_unload_entry=help_async_unload_entry,
+        ),
+    )
+    setup_test_component_platform(hass, DOMAIN, [entity], from_config_entry=True)
+    assert await hass.config_entries.async_setup(config_entry.entry_id)
+    await hass.async_block_till_done()
+
+    client = await hass_ws_client(hass)
+    await client.send_json_auto_id(
+        {"type": "vacuum/get_segments", "entity_id": entity.entity_id}
+    )
+    msg = await client.receive_json()
+    assert not msg["success"]
+    assert msg["error"]["code"] == ERR_NOT_SUPPORTED