Add group support to WLED main light

homeassistant/components/wled/light.py

@@ -16,6 +16,7 @@
 )
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
+from homeassistant.helpers.group import IntegrationSpecificGroup
 
 from .const import (
     ATTR_CCT,
@@ -61,11 +62,29 @@
     _attr_translation_key = "main"
     _attr_supported_features = LightEntityFeature.TRANSITION
     _attr_supported_color_modes = {ColorMode.BRIGHTNESS}
+    group: IntegrationSpecificGroup
 
     def __init__(self, coordinator: WLEDDataUpdateCoordinator) -> None:
         """Initialize WLED main light."""
         super().__init__(coordinator=coordinator)
         self._attr_unique_id = coordinator.data.info.mac_address
+        self.group = IntegrationSpecificGroup(self, [])
+        self._update_group_member()
+
+    def _update_group_member(self) -> None:
+        """Update group members based on current segments."""
+        segment_unique_ids = [
+            f"{self.coordinator.data.info.mac_address}_{segment.segment_id}"
+            for segment in sorted(
+                (
+                    segment
+                    for segment in self.coordinator.data.state.segments.values()
+                    if segment.segment_id is not None
+                ),
+                key=lambda segment: segment.segment_id,
+            )
+        ]
+        self.group.member_unique_ids = segment_unique_ids
 
     @property
     def brightness(self) -> int | None:
@@ -104,6 +123,12 @@
             on=True, brightness=kwargs.get(ATTR_BRIGHTNESS), transition=transition
         )
 
+    @callback
+    def _handle_coordinator_update(self) -> None:
+        """Update attributes when the coordinator updates."""
+        self._update_group_member()
+        super()._handle_coordinator_update()
+
 
 class WLEDSegmentLight(WLEDEntity, LightEntity):
     """Defines a WLED light based on a segment."""

tests/components/wled/snapshots/test_light.ambr

@@ -918,6 +918,10 @@
     ]),
     'area_id': None,
     'capabilities': dict({
+      'group_entities': list([
+        'light.wled_rgb_light',
+        'light.wled_rgb_light_segment_1',
+      ]),
       'supported_color_modes': list([
         <ColorMode.BRIGHTNESS: 'brightness'>,
       ]),
@@ -958,6 +962,10 @@
       'brightness': 128,
       'color_mode': <ColorMode.BRIGHTNESS: 'brightness'>,
       'friendly_name': 'WLED RGB Light Main',
+      'group_entities': list([
+        'light.wled_rgb_light',
+        'light.wled_rgb_light_segment_1',
+      ]),
       'supported_color_modes': list([
         <ColorMode.BRIGHTNESS: 'brightness'>,
       ]),

tests/components/wled/test_light.py

@@ -29,6 +29,7 @@
 )
 from homeassistant.const import (
     ATTR_ENTITY_ID,
+    ATTR_GROUP_ENTITIES,
     SERVICE_TURN_OFF,
     SERVICE_TURN_ON,
     STATE_OFF,
@@ -388,3 +389,52 @@ async def test_cct_light(hass: HomeAssistant, mock_wled: MagicMock) -> None:
         on=True,
         segment_id=0,
     )
+
+
+@pytest.mark.parametrize("device_fixture", ["rgb_single_segment"])
+async def test_main_light_group_updates_when_segments_change(
+    hass: HomeAssistant,
+    freezer: FrozenDateTimeFactory,
+    mock_wled: MagicMock,
+    init_integration: MockConfigEntry,
+) -> None:
+    """Test that the main light group field updates when segments are dynamically added or removed."""
+    single_segment_data = mock_wled.update.return_value
+    two_segment_data = WLEDDevice.from_dict(
+        await async_load_json_object_fixture(hass, "rgb.json", DOMAIN)
+    )
+
+    # Enable keep_main_light so the main light persists even with a single segment
+    hass.config_entries.async_update_entry(
+        init_integration, options={CONF_KEEP_MAIN_LIGHT: True}
+    )
+    await hass.config_entries.async_reload(init_integration.entry_id)
+    await hass.async_block_till_done()
+
+    # 1 segment: group should contain only segment 0
+    assert (state := hass.states.get("light.wled_rgb_light_main"))
+    assert state.state == STATE_ON
+    assert state.attributes[ATTR_GROUP_ENTITIES] == ["light.wled_rgb_light"]
+
+    # Add a second segment
+    mock_wled.update.return_value = two_segment_data
+    freezer.tick(SCAN_INTERVAL)
+    async_fire_time_changed(hass)
+    await hass.async_block_till_done()
+
+    # 2 segments: group should contain both
+    assert (state := hass.states.get("light.wled_rgb_light_main"))
+    assert state.attributes[ATTR_GROUP_ENTITIES] == [
+        "light.wled_rgb_light",
+        "light.wled_rgb_light_segment_1",
+    ]
+
+    # Remove the second segment
+    mock_wled.update.return_value = single_segment_data
+    freezer.tick(SCAN_INTERVAL)
+    async_fire_time_changed(hass)
+    await hass.async_block_till_done()
+
+    # Back to 1 segment: group should contain only segment 0 again
+    assert (state := hass.states.get("light.wled_rgb_light_main"))
+    assert state.attributes[ATTR_GROUP_ENTITIES] == ["light.wled_rgb_light"]