Update Shelly block valve platform to use entity description (#154068)

homeassistant/components/shelly/entity.py

@@ -79,6 +79,9 @@ def async_setup_block_attribute_entities(
                 continue
 
             # Filter out non-existing sensors and sensors without a value
+            if description.models and coordinator.model not in description.models:
+                continue
+
             if getattr(block, sensor_id, None) is None:
                 continue
 
@@ -301,6 +304,7 @@ class BlockEntityDescription(EntityDescription):
     available: Callable[[Block], bool] | None = None
     # Callable (settings, block), return true if entity should be removed
     removal_condition: Callable[[dict, Block], bool] | None = None
+    models: set[str] | None = None
 
 
 @dataclass(frozen=True, kw_only=True)

homeassistant/components/shelly/valve.py

@@ -27,7 +27,7 @@ from .entity import (
     async_setup_block_attribute_entities,
     async_setup_entry_rpc,
 )
-from .utils import async_remove_shelly_entity, get_device_entry_gen
+from .utils import get_device_entry_gen
 
 PARALLEL_UPDATES = 0
 
@@ -42,12 +42,15 @@ class RpcValveDescription(RpcEntityDescription, ValveEntityDescription):
     """Class to describe a RPC virtual valve."""
 
 
-GAS_VALVE = BlockValveDescription(
-    key="valve|valve",
-    name="Valve",
-    available=lambda block: block.valve not in ("failure", "checking"),
-    removal_condition=lambda _, block: block.valve in ("not_connected", "unknown"),
-)
+BLOCK_VALVES: dict[tuple[str, str], BlockValveDescription] = {
+    ("valve", "valve"): BlockValveDescription(
+        key="valve|valve",
+        name="Valve",
+        available=lambda block: block.valve not in ("failure", "checking"),
+        removal_condition=lambda _, block: block.valve in ("not_connected", "unknown"),
+        models={MODEL_GAS},
+    ),
+}
 
 
 class RpcShellyBaseWaterValve(ShellyRpcAttributeEntity, ValveEntity):
@@ -162,19 +165,15 @@ def async_setup_block_entry(
 ) -> None:
     """Set up valve for device."""
     coordinator = config_entry.runtime_data.block
-    assert coordinator and coordinator.device.blocks
+    assert coordinator
 
-    if coordinator.model == MODEL_GAS:
-        async_setup_block_attribute_entities(
-            hass,
-            async_add_entities,
-            coordinator,
-            {("valve", "valve"): GAS_VALVE},
-            BlockShellyValve,
-        )
-        # Remove deprecated switch entity for gas valve
-        unique_id = f"{coordinator.mac}-valve_0-valve"
-        async_remove_shelly_entity(hass, "switch", unique_id)
+    async_setup_block_attribute_entities(
+        hass,
+        async_add_entities,
+        coordinator,
+        BLOCK_VALVES,
+        BlockShellyValve,
+    )
 
 
 class BlockShellyValve(ShellyBlockAttributeEntity, ValveEntity):

tests/components/shelly/test_switch.py

@@ -4,7 +4,7 @@ from copy import deepcopy
 from datetime import timedelta
 from unittest.mock import AsyncMock, Mock
 
-from aioshelly.const import MODEL_1PM, MODEL_GAS, MODEL_MOTION
+from aioshelly.const import MODEL_1PM, MODEL_MOTION
 from aioshelly.exceptions import DeviceConnectionError, InvalidAuthError, RpcCallError
 from freezegun.api import FrozenDateTimeFactory
 import pytest
@@ -558,23 +558,6 @@ async def test_rpc_auth_error(
     assert flow["context"].get("entry_id") == entry.entry_id
 
 
-async def test_remove_gas_valve_switch(
-    hass: HomeAssistant,
-    mock_block_device: Mock,
-    entity_registry: EntityRegistry,
-) -> None:
-    """Test removing deprecated switch entity for Shelly Gas Valve."""
-    entity_id = register_entity(
-        hass,
-        SWITCH_DOMAIN,
-        "test_name_valve",
-        "valve_0-valve",
-    )
-    await init_integration(hass, 1, MODEL_GAS)
-
-    assert entity_registry.async_get(entity_id) is None
-
-
 async def test_wall_display_relay_mode(
     hass: HomeAssistant,
     mock_rpc_device: Mock,