Add Shelly support for valve entities (home-assistant#153348)

Author: thecode

homeassistant/components/shelly/const.py

@@ -308,3 +308,7 @@ class BLEScannerMode(StrEnum):
 MAX_SCRIPT_SIZE = 5120
 
 All_LIGHT_TYPES = ("cct", "light", "rgb", "rgbw")
+
+# Shelly-X specific models
+MODEL_NEO_WATER_VALVE = "NeoWaterValve"
+MODEL_FRANKEVER_WATER_VALVE = "WaterValve"

homeassistant/components/shelly/entity.py

@@ -186,6 +186,9 @@ def async_setup_rpc_attribute_entities(
 
         for key in key_instances:
             # Filter non-existing sensors
+            if description.models and coordinator.model not in description.models:
+                continue
+
             if description.role and description.role != coordinator.device.config[
                 key
             ].get("role", "generic"):
@@ -316,6 +319,7 @@ class RpcEntityDescription(EntityDescription):
     options_fn: Callable[[dict], list[str]] | None = None
     entity_class: Callable | None = None
     role: str | None = None
+    models: set[str] | None = None
 
 
 @dataclass(frozen=True)

homeassistant/components/shelly/valve.py

@@ -17,11 +17,15 @@
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
 
-from .coordinator import ShellyBlockCoordinator, ShellyConfigEntry
+from .const import MODEL_FRANKEVER_WATER_VALVE, MODEL_NEO_WATER_VALVE
+from .coordinator import ShellyBlockCoordinator, ShellyConfigEntry, ShellyRpcCoordinator
 from .entity import (
     BlockEntityDescription,
+    RpcEntityDescription,
     ShellyBlockAttributeEntity,
+    ShellyRpcAttributeEntity,
     async_setup_block_attribute_entities,
+    async_setup_entry_rpc,
 )
 from .utils import async_remove_shelly_entity, get_device_entry_gen
 
@@ -33,6 +37,11 @@ class BlockValveDescription(BlockEntityDescription, ValveEntityDescription):
     """Class to describe a BLOCK valve."""
 
 
+@dataclass(kw_only=True, frozen=True)
+class RpcValveDescription(RpcEntityDescription, ValveEntityDescription):
+    """Class to describe a RPC virtual valve."""
+
+
 GAS_VALVE = BlockValveDescription(
     key="valve|valve",
     name="Valve",
@@ -41,14 +50,108 @@ class BlockValveDescription(BlockEntityDescription, ValveEntityDescription):
 )
 
 
+class RpcShellyBaseWaterValve(ShellyRpcAttributeEntity, ValveEntity):
+    """Base Entity for RPC Shelly Water Valves."""
+
+    entity_description: RpcValveDescription
+    _attr_device_class = ValveDeviceClass.WATER
+    _id: int
+
+    def __init__(
+        self,
+        coordinator: ShellyRpcCoordinator,
+        key: str,
+        attribute: str,
+        description: RpcEntityDescription,
+    ) -> None:
+        """Initialize RPC water valve."""
+        super().__init__(coordinator, key, attribute, description)
+        self._attr_name = None  # Main device entity
+
+
+class RpcShellyWaterValve(RpcShellyBaseWaterValve):
+    """Entity that controls a valve on RPC Shelly Water Valve."""
+
+    _attr_supported_features = (
+        ValveEntityFeature.OPEN
+        | ValveEntityFeature.CLOSE
+        | ValveEntityFeature.SET_POSITION
+    )
+    _attr_reports_position = True
+
+    @property
+    def current_valve_position(self) -> int:
+        """Return current position of valve."""
+        return cast(int, self.attribute_value)
+
+    async def async_set_valve_position(self, position: int) -> None:
+        """Move the valve to a specific position."""
+        await self.coordinator.device.number_set(self._id, position)
+
+
+class RpcShellyNeoWaterValve(RpcShellyBaseWaterValve):
+    """Entity that controls a valve on RPC Shelly NEO Water Valve."""
+
+    _attr_supported_features = ValveEntityFeature.OPEN | ValveEntityFeature.CLOSE
+    _attr_reports_position = False
+
+    @property
+    def is_closed(self) -> bool | None:
+        """Return if the valve is closed or not."""
+        return not self.attribute_value
+
+    async def async_open_valve(self, **kwargs: Any) -> None:
+        """Open valve."""
+        await self.coordinator.device.boolean_set(self._id, True)
+
+    async def async_close_valve(self, **kwargs: Any) -> None:
+        """Close valve."""
+        await self.coordinator.device.boolean_set(self._id, False)
+
+
+RPC_VALVES: dict[str, RpcValveDescription] = {
+    "water_valve": RpcValveDescription(
+        key="number",
+        sub_key="value",
+        role="position",
+        entity_class=RpcShellyWaterValve,
+        models={MODEL_FRANKEVER_WATER_VALVE},
+    ),
+    "neo_water_valve": RpcValveDescription(
+        key="boolean",
+        sub_key="value",
+        role="state",
+        entity_class=RpcShellyNeoWaterValve,
+        models={MODEL_NEO_WATER_VALVE},
+    ),
+}
+
+
 async def async_setup_entry(
     hass: HomeAssistant,
     config_entry: ShellyConfigEntry,
     async_add_entities: AddConfigEntryEntitiesCallback,
 ) -> None:
     """Set up valves for device."""
     if get_device_entry_gen(config_entry) in BLOCK_GENERATIONS:
-        async_setup_block_entry(hass, config_entry, async_add_entities)
+        return async_setup_block_entry(hass, config_entry, async_add_entities)
+
+    return async_setup_rpc_entry(hass, config_entry, async_add_entities)
+
+
+@callback
+def async_setup_rpc_entry(
+    hass: HomeAssistant,
+    config_entry: ShellyConfigEntry,
+    async_add_entities: AddConfigEntryEntitiesCallback,
+) -> None:
+    """Set up entities for RPC device."""
+    coordinator = config_entry.runtime_data.rpc
+    assert coordinator
+
+    async_setup_entry_rpc(
+        hass, config_entry, async_add_entities, RPC_VALVES, RpcShellyWaterValve
+    )
 
 
 @callback

tests/components/shelly/test_valve.py

@@ -1,12 +1,27 @@
 """Tests for Shelly valve platform."""
 
+from copy import deepcopy
 from unittest.mock import Mock
 
 from aioshelly.const import MODEL_GAS
 import pytest
 
-from homeassistant.components.valve import DOMAIN as VALVE_DOMAIN, ValveState
-from homeassistant.const import ATTR_ENTITY_ID, SERVICE_CLOSE_VALVE, SERVICE_OPEN_VALVE
+from homeassistant.components.shelly.const import (
+    MODEL_FRANKEVER_WATER_VALVE,
+    MODEL_NEO_WATER_VALVE,
+)
+from homeassistant.components.valve import (
+    ATTR_CURRENT_POSITION,
+    ATTR_POSITION,
+    DOMAIN as VALVE_DOMAIN,
+    ValveState,
+)
+from homeassistant.const import (
+    ATTR_ENTITY_ID,
+    SERVICE_CLOSE_VALVE,
+    SERVICE_OPEN_VALVE,
+    SERVICE_SET_VALVE_POSITION,
+)
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_registry import EntityRegistry
 
@@ -64,3 +79,157 @@ async def test_block_device_gas_valve(
 
     assert (state := hass.states.get(entity_id))
     assert state.state == ValveState.CLOSED
+
+
+async def test_rpc_water_valve(
+    hass: HomeAssistant,
+    entity_registry: EntityRegistry,
+    mock_rpc_device: Mock,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """Test RPC device Shelly Water Valve."""
+    config = deepcopy(mock_rpc_device.config)
+    config["number:200"] = {
+        "name": "Position",
+        "min": 0,
+        "max": 100,
+        "meta": {"ui": {"step": 10, "view": "slider", "unit": "%"}},
+        "role": "position",
+    }
+    monkeypatch.setattr(mock_rpc_device, "config", config)
+
+    status = deepcopy(mock_rpc_device.status)
+    status["number:200"] = {"value": 0}
+    monkeypatch.setattr(mock_rpc_device, "status", status)
+
+    await init_integration(hass, 3, model=MODEL_FRANKEVER_WATER_VALVE)
+    entity_id = "valve.test_name"
+
+    assert (entry := entity_registry.async_get(entity_id))
+    assert entry.unique_id == "123456789ABC-number:200-water_valve"
+
+    assert (state := hass.states.get(entity_id))
+    assert state.state == ValveState.CLOSED
+
+    # Open valve
+    await hass.services.async_call(
+        VALVE_DOMAIN,
+        SERVICE_OPEN_VALVE,
+        {ATTR_ENTITY_ID: entity_id},
+        blocking=True,
+    )
+
+    mock_rpc_device.number_set.assert_called_once_with(200, 100)
+
+    status["number:200"] = {"value": 100}
+    monkeypatch.setattr(mock_rpc_device, "status", status)
+    mock_rpc_device.mock_update()
+    await hass.async_block_till_done()
+
+    assert (state := hass.states.get(entity_id))
+    assert state.state == ValveState.OPEN
+
+    # Close valve
+    mock_rpc_device.number_set.reset_mock()
+    await hass.services.async_call(
+        VALVE_DOMAIN,
+        SERVICE_CLOSE_VALVE,
+        {ATTR_ENTITY_ID: entity_id},
+        blocking=True,
+    )
+
+    mock_rpc_device.number_set.assert_called_once_with(200, 0)
+
+    status["number:200"] = {"value": 0}
+    monkeypatch.setattr(mock_rpc_device, "status", status)
+    mock_rpc_device.mock_update()
+    await hass.async_block_till_done()
+
+    assert (state := hass.states.get(entity_id))
+    assert state.state == ValveState.CLOSED
+
+    # Set valve position to 50%
+    mock_rpc_device.number_set.reset_mock()
+    await hass.services.async_call(
+        VALVE_DOMAIN,
+        SERVICE_SET_VALVE_POSITION,
+        {ATTR_ENTITY_ID: entity_id, ATTR_POSITION: 50},
+        blocking=True,
+    )
+
+    mock_rpc_device.number_set.assert_called_once_with(200, 50)
+
+    status["number:200"] = {"value": 50}
+    monkeypatch.setattr(mock_rpc_device, "status", status)
+    mock_rpc_device.mock_update()
+    await hass.async_block_till_done()
+
+    assert (state := hass.states.get(entity_id))
+    assert state.state == ValveState.OPEN
+    assert state.attributes.get(ATTR_CURRENT_POSITION) == 50
+
+
+async def test_rpc_neo_water_valve(
+    hass: HomeAssistant,
+    entity_registry: EntityRegistry,
+    mock_rpc_device: Mock,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """Test RPC device Shelly NEO Water Valve."""
+    config = deepcopy(mock_rpc_device.config)
+    config["boolean:200"] = {
+        "name": "State",
+        "meta": {"ui": {"view": "toggle"}},
+        "role": "state",
+    }
+    monkeypatch.setattr(mock_rpc_device, "config", config)
+
+    status = deepcopy(mock_rpc_device.status)
+    status["boolean:200"] = {"value": False}
+    monkeypatch.setattr(mock_rpc_device, "status", status)
+
+    await init_integration(hass, 3, model=MODEL_NEO_WATER_VALVE)
+    entity_id = "valve.test_name"
+
+    assert (entry := entity_registry.async_get(entity_id))
+    assert entry.unique_id == "123456789ABC-boolean:200-neo_water_valve"
+
+    assert (state := hass.states.get(entity_id))
+    assert state.state == ValveState.CLOSED
+
+    # Open valve
+    await hass.services.async_call(
+        VALVE_DOMAIN,
+        SERVICE_OPEN_VALVE,
+        {ATTR_ENTITY_ID: entity_id},
+        blocking=True,
+    )
+
+    mock_rpc_device.boolean_set.assert_called_once_with(200, True)
+
+    status["boolean:200"] = {"value": True}
+    monkeypatch.setattr(mock_rpc_device, "status", status)
+    mock_rpc_device.mock_update()
+    await hass.async_block_till_done()
+
+    assert (state := hass.states.get(entity_id))
+    assert state.state == ValveState.OPEN
+
+    # Close valve
+    mock_rpc_device.boolean_set.reset_mock()
+    await hass.services.async_call(
+        VALVE_DOMAIN,
+        SERVICE_CLOSE_VALVE,
+        {ATTR_ENTITY_ID: entity_id},
+        blocking=True,
+    )
+
+    mock_rpc_device.boolean_set.assert_called_once_with(200, False)
+
+    status["boolean:200"] = {"value": False}
+    monkeypatch.setattr(mock_rpc_device, "status", status)
+    mock_rpc_device.mock_update()
+    await hass.async_block_till_done()
+
+    assert (state := hass.states.get(entity_id))
+    assert state.state == ValveState.CLOSED