Add Ecovacs select entities (#108766)

Co-authored-by: Joost Lekkerkerker <joostlek@outlook.com>

homeassistant/components/ecovacs/__init__.py

@@ -26,6 +26,7 @@ CONFIG_SCHEMA = vol.Schema(
 
 PLATFORMS = [
     Platform.BINARY_SENSOR,
+    Platform.SELECT,
     Platform.SENSOR,
     Platform.VACUUM,
 ]

homeassistant/components/ecovacs/icons.json

@@ -45,6 +45,14 @@
       "total_stats_cleanings": {
         "default": "mdi:counter"
       }
+    },
+    "select": {
+      "water_amount": {
+        "default": "mdi:water"
+      },
+      "work_mode": {
+        "default": "mdi:cog"
+      }
     }
   }
 }

homeassistant/components/ecovacs/select.py

@@ -0,0 +1,101 @@
+"""Ecovacs select entity module."""
+from collections.abc import Callable
+from dataclasses import dataclass
+from typing import Any, Generic
+
+from deebot_client.capabilities import CapabilitySetTypes
+from deebot_client.device import Device
+from deebot_client.events import WaterInfoEvent, WorkModeEvent
+
+from homeassistant.components.select import SelectEntity, SelectEntityDescription
+from homeassistant.config_entries import ConfigEntry
+from homeassistant.const import EntityCategory
+from homeassistant.core import HomeAssistant
+from homeassistant.helpers.entity_platform import AddEntitiesCallback
+
+from .const import DOMAIN
+from .controller import EcovacsController
+from .entity import EcovacsCapabilityEntityDescription, EcovacsDescriptionEntity, EventT
+from .util import get_supported_entitites
+
+
+@dataclass(kw_only=True, frozen=True)
+class EcovacsSelectEntityDescription(
+    SelectEntityDescription,
+    EcovacsCapabilityEntityDescription,
+    Generic[EventT],
+):
+    """Ecovacs select entity description."""
+
+    current_option_fn: Callable[[EventT], str | None]
+    options_fn: Callable[[CapabilitySetTypes], list[str]]
+
+
+ENTITY_DESCRIPTIONS: tuple[EcovacsSelectEntityDescription, ...] = (
+    EcovacsSelectEntityDescription[WaterInfoEvent](
+        capability_fn=lambda caps: caps.water,
+        current_option_fn=lambda e: e.amount.display_name,
+        options_fn=lambda water: [amount.display_name for amount in water.types],
+        key="water_amount",
+        translation_key="water_amount",
+        entity_category=EntityCategory.CONFIG,
+    ),
+    EcovacsSelectEntityDescription[WorkModeEvent](
+        capability_fn=lambda caps: caps.clean.work_mode,
+        current_option_fn=lambda e: e.mode.display_name,
+        options_fn=lambda cap: [mode.display_name for mode in cap.types],
+        key="work_mode",
+        translation_key="work_mode",
+        entity_registry_enabled_default=False,
+        entity_category=EntityCategory.CONFIG,
+    ),
+)
+
+
+async def async_setup_entry(
+    hass: HomeAssistant,
+    config_entry: ConfigEntry,
+    async_add_entities: AddEntitiesCallback,
+) -> None:
+    """Add entities for passed config_entry in HA."""
+    controller: EcovacsController = hass.data[DOMAIN][config_entry.entry_id]
+    entities = get_supported_entitites(
+        controller, EcovacsSelectEntity, ENTITY_DESCRIPTIONS
+    )
+    if entities:
+        async_add_entities(entities)
+
+
+class EcovacsSelectEntity(
+    EcovacsDescriptionEntity[CapabilitySetTypes[EventT, str]],
+    SelectEntity,
+):
+    """Ecovacs select entity."""
+
+    _attr_current_option: str | None = None
+    entity_description: EcovacsSelectEntityDescription
+
+    def __init__(
+        self,
+        device: Device,
+        capability: CapabilitySetTypes[EventT, str],
+        entity_description: EcovacsSelectEntityDescription,
+        **kwargs: Any,
+    ) -> None:
+        """Initialize entity."""
+        super().__init__(device, capability, entity_description, **kwargs)
+        self._attr_options = entity_description.options_fn(capability)
+
+    async def async_added_to_hass(self) -> None:
+        """Set up the event listeners now that hass is ready."""
+        await super().async_added_to_hass()
+
+        async def on_event(event: EventT) -> None:
+            self._attr_current_option = self.entity_description.current_option_fn(event)
+            self.async_write_ha_state()
+
+        self._subscribe(self._capability.event, on_event)
+
+    async def async_select_option(self, option: str) -> None:
+        """Change the selected option."""
+        await self._device.execute_command(self._capability.set(option))

homeassistant/components/ecovacs/strings.json

@@ -67,6 +67,26 @@
         "name": "Total time cleaned"
       }
     },
+    "select": {
+      "water_amount": {
+        "name": "Water amount",
+        "state": {
+          "high": "High",
+          "low": "Low",
+          "medium": "Medium",
+          "ultrahigh": "Ultrahigh"
+        }
+      },
+      "work_mode": {
+        "name": "Work mode",
+        "state": {
+          "mop": "Mop",
+          "mop_after_vacuum": "Mop after vacuum",
+          "vacuum": "Vacuum",
+          "vacuum_and_mop": "Vacuum & mop"
+        }
+      }
+    },
     "vacuum": {
       "vacuum": {
         "state_attributes": {

tests/components/ecovacs/conftest.py

@@ -9,8 +9,10 @@ from deebot_client.exceptions import ApiError
 from deebot_client.models import Credentials
 import pytest
 
+from homeassistant.components.ecovacs import PLATFORMS
 from homeassistant.components.ecovacs.const import DOMAIN
 from homeassistant.components.ecovacs.controller import EcovacsController
+from homeassistant.const import Platform
 from homeassistant.core import HomeAssistant
 
 from .const import VALID_ENTRY_DATA
@@ -103,6 +105,12 @@ def mock_device_execute() -> AsyncMock:
         yield mock_device_execute
 
 
+@pytest.fixture
+def platforms() -> Platform | list[Platform]:
+    """Platforms, which should be loaded during the test."""
+    return PLATFORMS
+
+
 @pytest.fixture
 async def init_integration(
     hass: HomeAssistant,
@@ -110,13 +118,21 @@ async def init_integration(
     mock_authenticator: Mock,
     mock_mqtt_client: Mock,
     mock_device_execute: AsyncMock,
+    platforms: Platform | list[Platform],
 ) -> MockConfigEntry:
     """Set up the Ecovacs integration for testing."""
-    mock_config_entry.add_to_hass(hass)
+    if not isinstance(platforms, list):
+        platforms = [platforms]
 
-    await hass.config_entries.async_setup(mock_config_entry.entry_id)
+    with patch(
-    await hass.async_block_till_done()
+        "homeassistant.components.ecovacs.PLATFORMS",
-    return mock_config_entry
+        platforms,
+    ):
+        mock_config_entry.add_to_hass(hass)
+
+        await hass.config_entries.async_setup(mock_config_entry.entry_id)
+        await hass.async_block_till_done()
+        yield mock_config_entry
 
 
 @pytest.fixture

tests/components/ecovacs/snapshots/test_select.ambr

@@ -0,0 +1,57 @@
+# serializer version: 1
+# name: test_selects[yna5x1-entity_ids0][select.ozmo_950_water_amount:entity-registry]
+  EntityRegistryEntrySnapshot({
+    'aliases': set({
+    }),
+    'area_id': None,
+    'capabilities': dict({
+      'options': list([
+        'low',
+        'medium',
+        'high',
+        'ultrahigh',
+      ]),
+    }),
+    'config_entry_id': <ANY>,
+    'device_class': None,
+    'device_id': <ANY>,
+    'disabled_by': None,
+    'domain': 'select',
+    'entity_category': <EntityCategory.CONFIG: 'config'>,
+    'entity_id': 'select.ozmo_950_water_amount',
+    'has_entity_name': True,
+    'hidden_by': None,
+    'icon': None,
+    'id': <ANY>,
+    'name': None,
+    'options': dict({
+    }),
+    'original_device_class': None,
+    'original_icon': None,
+    'original_name': 'Water amount',
+    'platform': 'ecovacs',
+    'previous_unique_id': None,
+    'supported_features': 0,
+    'translation_key': 'water_amount',
+    'unique_id': 'E1234567890000000001_water_amount',
+    'unit_of_measurement': None,
+  })
+# ---
+# name: test_selects[yna5x1-entity_ids0][select.ozmo_950_water_amount:state]
+  StateSnapshot({
+    'attributes': ReadOnlyDict({
+      'friendly_name': 'Ozmo 950 Water amount',
+      'options': list([
+        'low',
+        'medium',
+        'high',
+        'ultrahigh',
+      ]),
+    }),
+    'context': <ANY>,
+    'entity_id': 'select.ozmo_950_water_amount',
+    'last_changed': <ANY>,
+    'last_updated': <ANY>,
+    'state': 'ultrahigh',
+  })
+# ---

tests/components/ecovacs/test_binary_sensor.py

@@ -1,25 +1,32 @@
 """Tests for Ecovacs binary sensors."""
 
-from deebot_client.event_bus import EventBus
 from deebot_client.events import WaterAmount, WaterInfoEvent
 import pytest
 from syrupy import SnapshotAssertion
 
+from homeassistant.components.ecovacs.const import DOMAIN
 from homeassistant.components.ecovacs.controller import EcovacsController
-from homeassistant.const import STATE_OFF, STATE_UNKNOWN
+from homeassistant.const import STATE_OFF, STATE_UNKNOWN, Platform
 from homeassistant.core import HomeAssistant
-from homeassistant.helpers import entity_registry as er
+from homeassistant.helpers import device_registry as dr, entity_registry as er
 
 from .util import notify_and_wait
 
 pytestmark = [pytest.mark.usefixtures("init_integration")]
 
 
+@pytest.fixture
+def platforms() -> Platform | list[Platform]:
+    """Platforms, which should be loaded during the test."""
+    return Platform.BINARY_SENSOR
+
+
 async def test_mop_attached(
     hass: HomeAssistant,
+    device_registry: dr.DeviceRegistry,
     entity_registry: er.EntityRegistry,
-    controller: EcovacsController,
     snapshot: SnapshotAssertion,
+    controller: EcovacsController,
 ) -> None:
     """Test mop_attached binary sensor."""
     entity_id = "binary_sensor.ozmo_950_mop_attached"
@@ -30,7 +37,12 @@ async def test_mop_attached(
     assert entity_entry == snapshot(name=f"{entity_id}-entity_entry")
     assert entity_entry.device_id
 
-    event_bus: EventBus = controller.devices[0].events
+    device = controller.devices[0]
+
+    assert (device_entry := device_registry.async_get(entity_entry.device_id))
+    assert device_entry.identifiers == {(DOMAIN, device.device_info.did)}
+
+    event_bus = device.events
     await notify_and_wait(
         hass, event_bus, WaterInfoEvent(WaterAmount.HIGH, mop_attached=True)
     )

tests/components/ecovacs/test_init.py

@@ -110,3 +110,21 @@ async def test_devices_in_dr(
             )
         )
         assert device_entry == snapshot(name=device.device_info.did)
+
+
+@pytest.mark.usefixtures("entity_registry_enabled_by_default", "init_integration")
+@pytest.mark.parametrize(
+    ("device_fixture", "entities"),
+    [
+        ("yna5x1", 16),
+    ],
+)
+async def test_all_entities_loaded(
+    hass: HomeAssistant,
+    device_fixture: str,
+    entities: int,
+) -> None:
+    """Test that all entities are loaded together."""
+    assert (
+        hass.states.async_entity_ids_count() == entities
+    ), f"loaded entities for {device_fixture}: {hass.states.async_entity_ids()}"

tests/components/ecovacs/test_select.py

@@ -0,0 +1,115 @@
+"""Tests for Ecovacs select entities."""
+
+from deebot_client.command import Command
+from deebot_client.commands.json import SetWaterInfo
+from deebot_client.event_bus import EventBus
+from deebot_client.events import WaterAmount, WaterInfoEvent
+import pytest
+from syrupy import SnapshotAssertion
+
+from homeassistant.components import select
+from homeassistant.components.ecovacs.const import DOMAIN
+from homeassistant.components.ecovacs.controller import EcovacsController
+from homeassistant.const import (
+    ATTR_ENTITY_ID,
+    ATTR_OPTION,
+    SERVICE_SELECT_OPTION,
+    STATE_UNKNOWN,
+    Platform,
+)
+from homeassistant.core import HomeAssistant
+from homeassistant.helpers import device_registry as dr, entity_registry as er
+
+from .util import block_till_done
+
+pytestmark = [pytest.mark.usefixtures("init_integration")]
+
+
+@pytest.fixture
+def platforms() -> Platform | list[Platform]:
+    """Platforms, which should be loaded during the test."""
+    return Platform.SELECT
+
+
+async def notify_events(hass: HomeAssistant, event_bus: EventBus):
+    """Notify events."""
+    event_bus.notify(WaterInfoEvent(WaterAmount.ULTRAHIGH))
+    await block_till_done(hass, event_bus)
+
+
+@pytest.mark.usefixtures("entity_registry_enabled_by_default")
+@pytest.mark.parametrize(
+    ("device_fixture", "entity_ids"),
+    [
+        (
+            "yna5x1",
+            [
+                "select.ozmo_950_water_amount",
+            ],
+        ),
+    ],
+)
+async def test_selects(
+    hass: HomeAssistant,
+    device_registry: dr.DeviceRegistry,
+    entity_registry: er.EntityRegistry,
+    snapshot: SnapshotAssertion,
+    controller: EcovacsController,
+    entity_ids: list[str],
+) -> None:
+    """Test that select entity snapshots match."""
+    assert entity_ids == sorted(hass.states.async_entity_ids())
+    for entity_id in entity_ids:
+        assert (state := hass.states.get(entity_id)), f"State of {entity_id} is missing"
+        assert state.state == STATE_UNKNOWN
+
+    device = controller.devices[0]
+    await notify_events(hass, device.events)
+    for entity_id in entity_ids:
+        assert (state := hass.states.get(entity_id)), f"State of {entity_id} is missing"
+        assert snapshot(name=f"{entity_id}:state") == state
+
+        assert (entity_entry := entity_registry.async_get(state.entity_id))
+        assert snapshot(name=f"{entity_id}:entity-registry") == entity_entry
+
+        assert entity_entry.device_id
+        assert (device_entry := device_registry.async_get(entity_entry.device_id))
+        assert device_entry.identifiers == {(DOMAIN, device.device_info.did)}
+
+
+@pytest.mark.usefixtures("entity_registry_enabled_by_default")
+@pytest.mark.parametrize(
+    ("device_fixture", "entity_id", "current_state", "set_state", "command"),
+    [
+        (
+            "yna5x1",
+            "select.ozmo_950_water_amount",
+            "ultrahigh",
+            "low",
+            SetWaterInfo(WaterAmount.LOW),
+        ),
+    ],
+)
+async def test_selects_change(
+    hass: HomeAssistant,
+    controller: EcovacsController,
+    entity_id: list[str],
+    current_state: str,
+    set_state: str,
+    command: Command,
+) -> None:
+    """Test that changing select entities works."""
+    device = controller.devices[0]
+    await notify_events(hass, device.events)
+
+    assert (state := hass.states.get(entity_id)), f"State of {entity_id} is missing"
+    assert state.state == current_state
+
+    device._execute_command.reset_mock()
+    await hass.services.async_call(
+        select.DOMAIN,
+        SERVICE_SELECT_OPTION,
+        {ATTR_ENTITY_ID: entity_id, ATTR_OPTION: set_state},
+        blocking=True,
+    )
+    device._execute_command.assert_called_with(command)

tests/components/ecovacs/test_sensor.py

@@ -13,16 +13,23 @@ from deebot_client.events import (
 import pytest
 from syrupy import SnapshotAssertion
 
+from homeassistant.components.ecovacs.const import DOMAIN
 from homeassistant.components.ecovacs.controller import EcovacsController
 from homeassistant.const import STATE_UNKNOWN, Platform
 from homeassistant.core import HomeAssistant
-from homeassistant.helpers import entity_registry as er
+from homeassistant.helpers import device_registry as dr, entity_registry as er
 
 from .util import block_till_done
 
 pytestmark = [pytest.mark.usefixtures("init_integration")]
 
 
+@pytest.fixture
+def platforms() -> Platform | list[Platform]:
+    """Platforms, which should be loaded during the test."""
+    return Platform.SENSOR
+
+
 async def notify_events(hass: HomeAssistant, event_bus: EventBus):
     """Notify events."""
     event_bus.notify(StatsEvent(10, 300, "spotArea"))
@@ -65,18 +72,20 @@ async def notify_events(hass: HomeAssistant, event_bus: EventBus):
 )
 async def test_sensors(
     hass: HomeAssistant,
-    controller: EcovacsController,
+    device_registry: dr.DeviceRegistry,
     entity_registry: er.EntityRegistry,
     snapshot: SnapshotAssertion,
+    controller: EcovacsController,
     entity_ids: list[str],
 ) -> None:
     """Test that sensor entity snapshots match."""
-    assert entity_ids == sorted(hass.states.async_entity_ids(Platform.SENSOR))
+    assert entity_ids == sorted(hass.states.async_entity_ids())
     for entity_id in entity_ids:
         assert (state := hass.states.get(entity_id)), f"State of {entity_id} is missing"
         assert state.state == STATE_UNKNOWN
 
-    await notify_events(hass, controller.devices[0].events)
+    device = controller.devices[0]
+    await notify_events(hass, device.events)
     for entity_id in entity_ids:
         assert (state := hass.states.get(entity_id)), f"State of {entity_id} is missing"
         assert snapshot(name=f"{entity_id}:state") == state
@@ -85,6 +94,8 @@ async def test_sensors(
         assert snapshot(name=f"{entity_id}:entity-registry") == entity_entry
 
         assert entity_entry.device_id
+        assert (device_entry := device_registry.async_get(entity_entry.device_id))
+        assert device_entry.identifiers == {(DOMAIN, device.device_info.did)}
 
 
 @pytest.mark.parametrize(