Use Entity Description in Shelly light platform (home-assistant#154102)

Author: thecode

homeassistant/components/shelly/light.py

@@ -40,14 +40,15 @@
 )
 from .coordinator import ShellyBlockCoordinator, ShellyConfigEntry, ShellyRpcCoordinator
 from .entity import (
+    BlockEntityDescription,
     RpcEntityDescription,
-    ShellyBlockEntity,
+    ShellyBlockAttributeEntity,
     ShellyRpcAttributeEntity,
+    async_setup_entry_attribute_entities,
     async_setup_entry_rpc,
 )
 from .utils import (
     async_remove_orphaned_entities,
-    async_remove_shelly_entity,
     brightness_to_percentage,
     get_device_entry_gen,
     is_block_channel_type_light,
@@ -58,6 +59,24 @@
 PARALLEL_UPDATES = 0
 
 
+@dataclass(frozen=True, kw_only=True)
+class BlockLightDescription(BlockEntityDescription, LightEntityDescription):
+    """Description for a Shelly BLOCK light entity."""
+
+
+BLOCK_LIGHTS = {
+    ("light", "output"): BlockLightDescription(
+        key="light|output",
+    ),
+    ("relay", "output"): BlockLightDescription(
+        key="relay|output",
+        removal_condition=lambda settings, block: not is_block_channel_type_light(
+            settings, block
+        ),
+    ),
+}
+
+
 async def async_setup_entry(
     hass: HomeAssistant,
     config_entry: ShellyConfigEntry,
@@ -79,36 +98,29 @@ def async_setup_block_entry(
     """Set up entities for block device."""
     coordinator = config_entry.runtime_data.block
     assert coordinator
-    blocks = []
-    assert coordinator.device.blocks
-    for block in coordinator.device.blocks:
-        if block.type == "light":
-            blocks.append(block)
-        elif block.type == "relay" and block.channel is not None:
-            if not is_block_channel_type_light(
-                coordinator.device.settings, int(block.channel)
-            ):
-                continue
 
-            blocks.append(block)
-            unique_id = f"{coordinator.mac}-{block.type}_{block.channel}"
-            async_remove_shelly_entity(hass, "switch", unique_id)
-
-    if not blocks:
-        return
-
-    async_add_entities(BlockShellyLight(coordinator, block) for block in blocks)
+    async_setup_entry_attribute_entities(
+        hass, config_entry, async_add_entities, BLOCK_LIGHTS, BlockShellyLight
+    )
 
 
-class BlockShellyLight(ShellyBlockEntity, LightEntity):
+class BlockShellyLight(ShellyBlockAttributeEntity, LightEntity):
     """Entity that controls a light on block based Shelly devices."""
 
+    entity_description: BlockLightDescription
     _attr_supported_color_modes: set[str]
 
-    def __init__(self, coordinator: ShellyBlockCoordinator, block: Block) -> None:
-        """Initialize light."""
-        super().__init__(coordinator, block)
+    def __init__(
+        self,
+        coordinator: ShellyBlockCoordinator,
+        block: Block,
+        attribute: str,
+        description: BlockLightDescription,
+    ) -> None:
+        """Initialize block light."""
+        super().__init__(coordinator, block, attribute, description)
         self.control_result: dict[str, Any] | None = None
+        self._attr_unique_id: str = f"{coordinator.mac}-{block.description}"
         self._attr_supported_color_modes = set()
         self._attr_min_color_temp_kelvin = KELVIN_MIN_VALUE_WHITE
         self._attr_max_color_temp_kelvin = KELVIN_MAX_VALUE
@@ -347,7 +359,7 @@ def _update_callback(self) -> None:
 
 @dataclass(frozen=True, kw_only=True)
 class RpcLightDescription(RpcEntityDescription, LightEntityDescription):
-    """Description for a Shelly RPC number entity."""
+    """Description for a Shelly RPC light entity."""
 
 
 class RpcShellyLightBase(ShellyRpcAttributeEntity, LightEntity):

homeassistant/components/shelly/utils.py

@@ -211,10 +211,7 @@ def is_block_exclude_from_relay(settings: dict[str, Any], block: Block) -> bool:
     if settings.get("mode") == "roller":
         return True
 
-    if TYPE_CHECKING:
-        assert block.channel is not None
-
-    return is_block_channel_type_light(settings, int(block.channel))
+    return is_block_channel_type_light(settings, block)
 
 
 def get_device_uptime(uptime: float, last_uptime: datetime | None) -> datetime:
@@ -501,9 +498,12 @@ def is_rpc_momentary_input(
     return cast(bool, config[key]["type"] == "button")
 
 
-def is_block_channel_type_light(settings: dict[str, Any], channel: int) -> bool:
+def is_block_channel_type_light(settings: dict[str, Any], block: Block) -> bool:
     """Return true if block channel appliance type is set to light."""
-    app_type = settings["relays"][channel].get("appliance_type")
+    if TYPE_CHECKING:
+        assert block.channel is not None
+
+    app_type = settings["relays"][int(block.channel)].get("appliance_type")
     return app_type is not None and app_type.lower().startswith("light")
 
 

tests/components/shelly/conftest.py

@@ -127,7 +127,11 @@ def mock_white_light_set_state(
         ),
     ),
     Mock(
-        sensor_ids={"mode": "color", "effect": 0},
+        sensor_ids={
+            "output": mock_light_set_state()["ison"],
+            "mode": "color",
+            "effect": 0,
+        },
         channel="0",
         output=mock_light_set_state()["ison"],
         colorTemp=mock_light_set_state()["temp"],