Add support for Shelly `enum` virtual component (#121997)

* Support enum virtual component

* Add tests

* Cleaning

* Improve test for select

* Use values

* Update tests

* Use the option title for sensor

---------

Co-authored-by: Maciej Bieniek <478555+bieniu@users.noreply.github.com>

homeassistant/components/shelly/__init__.py

@@ -62,6 +62,7 @@ PLATFORMS: Final = [
     Platform.EVENT,
     Platform.LIGHT,
     Platform.NUMBER,
+    Platform.SELECT,
     Platform.SENSOR,
     Platform.SWITCH,
     Platform.TEXT,

homeassistant/components/shelly/const.py

@@ -244,7 +244,8 @@ SHELLY_PLUS_RGBW_CHANNELS = 4
 VIRTUAL_COMPONENTS_MAP = {
     "binary_sensor": {"types": ["boolean"], "modes": ["label"]},
     "number": {"types": ["number"], "modes": ["field", "slider"]},
-    "sensor": {"types": ["number", "text"], "modes": ["label"]},
+    "select": {"types": ["enum"], "modes": ["dropdown"]},
+    "sensor": {"types": ["enum", "number", "text"], "modes": ["label"]},
     "switch": {"types": ["boolean"], "modes": ["toggle"]},
     "text": {"types": ["text"], "modes": ["field"]},
 }

homeassistant/components/shelly/entity.py

@@ -292,6 +292,7 @@ class RpcEntityDescription(EntityDescription):
     use_polling_coordinator: bool = False
     supported: Callable = lambda _: False
     unit: Callable[[dict], str | None] | None = None
+    options_fn: Callable[[dict], list[str]] | None = None
 
 
 @dataclass(frozen=True)
@@ -514,6 +515,19 @@ class ShellyRpcAttributeEntity(ShellyRpcEntity, Entity):
                 coordinator.device.config[key]
             )
 
+        self.option_map: dict[str, str] = {}
+        self.reversed_option_map: dict[str, str] = {}
+        if "enum" in key:
+            titles = self.coordinator.device.config[key]["meta"]["ui"]["titles"]
+            options = self.coordinator.device.config[key]["options"]
+            self.option_map = {
+                opt: (titles[opt] if titles.get(opt) is not None else opt)
+                for opt in options
+            }
+            self.reversed_option_map = {
+                tit: opt for opt, tit in self.option_map.items()
+            }
+
     @property
     def sub_status(self) -> Any:
         """Device status by entity key."""

homeassistant/components/shelly/select.py

@@ -0,0 +1,103 @@
+"""Select for Shelly."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Final
+
+from aioshelly.const import RPC_GENERATIONS
+
+from homeassistant.components.select import (
+    DOMAIN as SELECT_PLATFORM,
+    SelectEntity,
+    SelectEntityDescription,
+)
+from homeassistant.core import HomeAssistant
+from homeassistant.helpers.entity_platform import AddEntitiesCallback
+
+from .coordinator import ShellyConfigEntry, ShellyRpcCoordinator
+from .entity import (
+    RpcEntityDescription,
+    ShellyRpcAttributeEntity,
+    async_setup_entry_rpc,
+)
+from .utils import (
+    async_remove_orphaned_virtual_entities,
+    get_device_entry_gen,
+    get_virtual_component_ids,
+)
+
+
+@dataclass(frozen=True, kw_only=True)
+class RpcSelectDescription(RpcEntityDescription, SelectEntityDescription):
+    """Class to describe a RPC select entity."""
+
+
+RPC_SELECT_ENTITIES: Final = {
+    "enum": RpcSelectDescription(
+        key="enum",
+        sub_key="value",
+        has_entity_name=True,
+    ),
+}
+
+
+async def async_setup_entry(
+    hass: HomeAssistant,
+    config_entry: ShellyConfigEntry,
+    async_add_entities: AddEntitiesCallback,
+) -> None:
+    """Set up selectors for device."""
+    if get_device_entry_gen(config_entry) in RPC_GENERATIONS:
+        coordinator = config_entry.runtime_data.rpc
+        assert coordinator
+
+        async_setup_entry_rpc(
+            hass, config_entry, async_add_entities, RPC_SELECT_ENTITIES, RpcSelect
+        )
+
+        # the user can remove virtual components from the device configuration, so
+        # we need to remove orphaned entities
+        virtual_text_ids = get_virtual_component_ids(
+            coordinator.device.config, SELECT_PLATFORM
+        )
+        async_remove_orphaned_virtual_entities(
+            hass,
+            config_entry.entry_id,
+            coordinator.mac,
+            SELECT_PLATFORM,
+            "enum",
+            virtual_text_ids,
+        )
+
+
+class RpcSelect(ShellyRpcAttributeEntity, SelectEntity):
+    """Represent a RPC select entity."""
+
+    entity_description: RpcSelectDescription
+
+    def __init__(
+        self,
+        coordinator: ShellyRpcCoordinator,
+        key: str,
+        attribute: str,
+        description: RpcSelectDescription,
+    ) -> None:
+        """Initialize select."""
+        super().__init__(coordinator, key, attribute, description)
+
+        self._attr_options = list(self.option_map.values())
+
+    @property
+    def current_option(self) -> str | None:
+        """Return the selected entity option to represent the entity state."""
+        if not isinstance(self.attribute_value, str):
+            return None
+
+        return self.option_map[self.attribute_value]
+
+    async def async_select_option(self, option: str) -> None:
+        """Change the value."""
+        await self.call_rpc(
+            "Enum.Set", {"id": self._id, "value": self.reversed_option_map[option]}
+        )

homeassistant/components/shelly/sensor.py

@@ -1032,6 +1032,13 @@ RPC_SENSORS: Final = {
         if config["meta"]["ui"]["unit"]
         else None,
     ),
+    "enum": RpcSensorDescription(
+        key="enum",
+        sub_key="value",
+        has_entity_name=True,
+        options_fn=lambda config: config["options"],
+        device_class=SensorDeviceClass.ENUM,
+    ),
 }
 
 
@@ -1060,7 +1067,7 @@ async def async_setup_entry(
 
             # the user can remove virtual components from the device configuration, so
             # we need to remove orphaned entities
-            for component in ("text", "number"):
+            for component in ("enum", "number", "text"):
                 virtual_component_ids = get_virtual_component_ids(
                     coordinator.device.config, SENSOR_PLATFORM
                 )
@@ -1134,10 +1141,29 @@ class RpcSensor(ShellyRpcAttributeEntity, SensorEntity):
 
     entity_description: RpcSensorDescription
 
+    def __init__(
+        self,
+        coordinator: ShellyRpcCoordinator,
+        key: str,
+        attribute: str,
+        description: RpcSensorDescription,
+    ) -> None:
+        """Initialize select."""
+        super().__init__(coordinator, key, attribute, description)
+
+        if self.option_map:
+            self._attr_options = list(self.option_map.values())
+
     @property
     def native_value(self) -> StateType:
         """Return value of sensor."""
-        return self.attribute_value
+        if not self.option_map:
+            return self.attribute_value
+
+        if not isinstance(self.attribute_value, str):
+            return None
+
+        return self.option_map[self.attribute_value]
 
 
 class BlockSleepingSensor(ShellySleepingBlockAttributeEntity, RestoreSensor):

homeassistant/components/shelly/utils.py

@@ -323,7 +323,7 @@ def get_rpc_channel_name(device: RpcDevice, key: str) -> str:
             return f"{device_name} {key.replace(':', '_')}"
         if key.startswith("em1"):
             return f"{device_name} EM{key.split(':')[-1]}"
-        if key.startswith(("boolean:", "number:", "text:")):
+        if key.startswith(("boolean:", "enum:", "number:", "text:")):
             return key.replace(":", " ").title()
         return device_name
 

tests/components/shelly/test_select.py

@@ -0,0 +1,151 @@
+"""Tests for Shelly select platform."""
+
+from copy import deepcopy
+from unittest.mock import Mock
+
+import pytest
+
+from homeassistant.components.select import (
+    ATTR_OPTION,
+    ATTR_OPTIONS,
+    DOMAIN as SELECT_PLATFORM,
+    SERVICE_SELECT_OPTION,
+)
+from homeassistant.const import ATTR_ENTITY_ID, STATE_UNKNOWN
+from homeassistant.core import HomeAssistant
+from homeassistant.helpers.device_registry import DeviceRegistry
+from homeassistant.helpers.entity_registry import EntityRegistry
+
+from . import init_integration, register_device, register_entity
+
+
+@pytest.mark.parametrize(
+    ("name", "entity_id", "value", "expected_state"),
+    [
+        ("Virtual enum", "select.test_name_virtual_enum", "option 1", "Title 1"),
+        (None, "select.test_name_enum_203", None, STATE_UNKNOWN),
+    ],
+)
+async def test_rpc_device_virtual_enum(
+    hass: HomeAssistant,
+    entity_registry: EntityRegistry,
+    mock_rpc_device: Mock,
+    monkeypatch: pytest.MonkeyPatch,
+    name: str | None,
+    entity_id: str,
+    value: str | None,
+    expected_state: str,
+) -> None:
+    """Test a virtual enum for RPC device."""
+    config = deepcopy(mock_rpc_device.config)
+    config["enum:203"] = {
+        "name": name,
+        "options": ["option 1", "option 2", "option 3"],
+        "meta": {
+            "ui": {
+                "view": "dropdown",
+                "titles": {"option 1": "Title 1", "option 2": None},
+            }
+        },
+    }
+    monkeypatch.setattr(mock_rpc_device, "config", config)
+
+    status = deepcopy(mock_rpc_device.status)
+    status["enum:203"] = {"value": value}
+    monkeypatch.setattr(mock_rpc_device, "status", status)
+
+    await init_integration(hass, 3)
+
+    state = hass.states.get(entity_id)
+    assert state
+    assert state.state == expected_state
+    assert state.attributes.get(ATTR_OPTIONS) == [
+        "Title 1",
+        "option 2",
+        "option 3",
+    ]
+
+    entry = entity_registry.async_get(entity_id)
+    assert entry
+    assert entry.unique_id == "123456789ABC-enum:203-enum"
+
+    monkeypatch.setitem(mock_rpc_device.status["enum:203"], "value", "option 2")
+    mock_rpc_device.mock_update()
+    assert hass.states.get(entity_id).state == "option 2"
+
+    monkeypatch.setitem(mock_rpc_device.status["enum:203"], "value", "option 1")
+    await hass.services.async_call(
+        SELECT_PLATFORM,
+        SERVICE_SELECT_OPTION,
+        {ATTR_ENTITY_ID: entity_id, ATTR_OPTION: "Title 1"},
+        blocking=True,
+    )
+    # 'Title 1' corresponds to 'option 1'
+    assert mock_rpc_device.call_rpc.call_args[0][1] == {"id": 203, "value": "option 1"}
+    mock_rpc_device.mock_update()
+    assert hass.states.get(entity_id).state == "Title 1"
+
+
+async def test_rpc_remove_virtual_enum_when_mode_label(
+    hass: HomeAssistant,
+    entity_registry: EntityRegistry,
+    device_registry: DeviceRegistry,
+    mock_rpc_device: Mock,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """Test if the virtual enum will be removed if the mode has been changed to a label."""
+    config = deepcopy(mock_rpc_device.config)
+    config["enum:200"] = {
+        "name": None,
+        "options": ["one", "two"],
+        "meta": {
+            "ui": {"view": "label", "titles": {"one": "Title 1", "two": "Title 2"}}
+        },
+    }
+    monkeypatch.setattr(mock_rpc_device, "config", config)
+
+    status = deepcopy(mock_rpc_device.status)
+    status["enum:200"] = {"value": "one"}
+    monkeypatch.setattr(mock_rpc_device, "status", status)
+
+    config_entry = await init_integration(hass, 3, skip_setup=True)
+    device_entry = register_device(device_registry, config_entry)
+    entity_id = register_entity(
+        hass,
+        SELECT_PLATFORM,
+        "test_name_enum_200",
+        "enum:200-enum",
+        config_entry,
+        device_id=device_entry.id,
+    )
+
+    await hass.config_entries.async_setup(config_entry.entry_id)
+    await hass.async_block_till_done()
+
+    entry = entity_registry.async_get(entity_id)
+    assert not entry
+
+
+async def test_rpc_remove_virtual_enum_when_orphaned(
+    hass: HomeAssistant,
+    entity_registry: EntityRegistry,
+    device_registry: DeviceRegistry,
+    mock_rpc_device: Mock,
+) -> None:
+    """Check whether the virtual enum will be removed if it has been removed from the device configuration."""
+    config_entry = await init_integration(hass, 3, skip_setup=True)
+    device_entry = register_device(device_registry, config_entry)
+    entity_id = register_entity(
+        hass,
+        SELECT_PLATFORM,
+        "test_name_enum_200",
+        "enum:200-enum",
+        config_entry,
+        device_id=device_entry.id,
+    )
+
+    await hass.config_entries.async_setup(config_entry.entry_id)
+    await hass.async_block_till_done()
+
+    entry = entity_registry.async_get(entity_id)
+    assert not entry

tests/components/shelly/test_sensor.py

@@ -11,6 +11,7 @@ from homeassistant.components.homeassistant import (
     SERVICE_UPDATE_ENTITY,
 )
 from homeassistant.components.sensor import (
+    ATTR_OPTIONS,
     ATTR_STATE_CLASS,
     DOMAIN as SENSOR_DOMAIN,
     SensorDeviceClass,
@@ -1066,3 +1067,123 @@ async def test_rpc_remove_number_virtual_sensor_when_orphaned(
 
     entry = entity_registry.async_get(entity_id)
     assert not entry
+
+
+@pytest.mark.parametrize(
+    ("name", "entity_id", "value", "expected_state"),
+    [
+        (
+            "Virtual enum sensor",
+            "sensor.test_name_virtual_enum_sensor",
+            "one",
+            "Title 1",
+        ),
+        (None, "sensor.test_name_enum_203", None, STATE_UNKNOWN),
+    ],
+)
+async def test_rpc_device_virtual_enum_sensor(
+    hass: HomeAssistant,
+    entity_registry: EntityRegistry,
+    mock_rpc_device: Mock,
+    monkeypatch: pytest.MonkeyPatch,
+    name: str | None,
+    entity_id: str,
+    value: str | None,
+    expected_state: str,
+) -> None:
+    """Test a virtual enum sensor for RPC device."""
+    config = deepcopy(mock_rpc_device.config)
+    config["enum:203"] = {
+        "name": name,
+        "options": ["one", "two", "three"],
+        "meta": {"ui": {"view": "label", "titles": {"one": "Title 1", "two": None}}},
+    }
+    monkeypatch.setattr(mock_rpc_device, "config", config)
+
+    status = deepcopy(mock_rpc_device.status)
+    status["enum:203"] = {"value": value}
+    monkeypatch.setattr(mock_rpc_device, "status", status)
+
+    await init_integration(hass, 3)
+
+    state = hass.states.get(entity_id)
+    assert state
+    assert state.state == expected_state
+    assert state.attributes.get(ATTR_DEVICE_CLASS) == SensorDeviceClass.ENUM
+    assert state.attributes.get(ATTR_OPTIONS) == ["Title 1", "two", "three"]
+
+    entry = entity_registry.async_get(entity_id)
+    assert entry
+    assert entry.unique_id == "123456789ABC-enum:203-enum"
+
+    monkeypatch.setitem(mock_rpc_device.status["enum:203"], "value", "two")
+    mock_rpc_device.mock_update()
+    assert hass.states.get(entity_id).state == "two"
+
+
+async def test_rpc_remove_enum_virtual_sensor_when_mode_dropdown(
+    hass: HomeAssistant,
+    entity_registry: EntityRegistry,
+    device_registry: DeviceRegistry,
+    mock_rpc_device: Mock,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """Test if the virtual enum sensor will be removed if the mode has been changed to a dropdown."""
+    config = deepcopy(mock_rpc_device.config)
+    config["enum:200"] = {
+        "name": None,
+        "options": ["option 1", "option 2", "option 3"],
+        "meta": {
+            "ui": {
+                "view": "dropdown",
+                "titles": {"option 1": "Title 1", "option 2": None},
+            }
+        },
+    }
+    monkeypatch.setattr(mock_rpc_device, "config", config)
+
+    status = deepcopy(mock_rpc_device.status)
+    status["enum:200"] = {"value": "option 2"}
+    monkeypatch.setattr(mock_rpc_device, "status", status)
+
+    config_entry = await init_integration(hass, 3, skip_setup=True)
+    device_entry = register_device(device_registry, config_entry)
+    entity_id = register_entity(
+        hass,
+        SENSOR_DOMAIN,
+        "test_name_enum_200",
+        "enum:200-enum",
+        config_entry,
+        device_id=device_entry.id,
+    )
+
+    await hass.config_entries.async_setup(config_entry.entry_id)
+    await hass.async_block_till_done()
+
+    entry = entity_registry.async_get(entity_id)
+    assert not entry
+
+
+async def test_rpc_remove_enum_virtual_sensor_when_orphaned(
+    hass: HomeAssistant,
+    entity_registry: EntityRegistry,
+    device_registry: DeviceRegistry,
+    mock_rpc_device: Mock,
+) -> None:
+    """Check whether the virtual enum sensor will be removed if it has been removed from the device configuration."""
+    config_entry = await init_integration(hass, 3, skip_setup=True)
+    device_entry = register_device(device_registry, config_entry)
+    entity_id = register_entity(
+        hass,
+        SENSOR_DOMAIN,
+        "test_name_enum_200",
+        "enum:200-enum",
+        config_entry,
+        device_id=device_entry.id,
+    )
+
+    await hass.config_entries.async_setup(config_entry.entry_id)
+    await hass.async_block_till_done()
+
+    entry = entity_registry.async_get(entity_id)
+    assert not entry