Add group support to WLED main light

homeassistant/components/wled/coordinator.py

@@ -81,6 +81,15 @@ def has_main_light(self) -> bool:
             self.data is not None and len(self.data.state.segments) > 1
         )
 
+    @property
+    def segment_ids(self) -> set[int]:
+        """Return the set of segment IDs."""
+        return {
+            segment.segment_id
+            for segment in self.data.state.segments.values()
+            if segment.segment_id is not None
+        }
+
     @callback
     def _use_websocket(self) -> None:
         """Use WebSocket for updates, instead of polling."""

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,23 @@ class WLEDMainLight(WLEDEntity, LightEntity):
     _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_id}"
+            for segment_id in sorted(self.coordinator.segment_ids)
+        ]
+        if segment_unique_ids != self.group.member_unique_ids:
+            self.group.member_unique_ids = segment_unique_ids
 
     @property
     def brightness(self) -> int | None:
@@ -104,6 +117,12 @@ async def async_turn_on(self, **kwargs: Any) -> None:
             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."""
@@ -280,11 +299,7 @@ def async_update_segments(
     async_add_entities: AddConfigEntryEntitiesCallback,
 ) -> None:
     """Update segments."""
-    segment_ids = {
-        light.segment_id
-        for light in coordinator.data.state.segments.values()
-        if light.segment_id is not None
-    }
+    segment_ids = coordinator.segment_ids
     new_entities: list[WLEDMainLight | WLEDSegmentLight] = []
 
     # More than 1 segment now? No main? Add main controls

homeassistant/components/wled/number.py

@@ -127,16 +127,10 @@ def async_update_segments(
     async_add_entities: AddConfigEntryEntitiesCallback,
 ) -> None:
     """Update segments."""
-    segment_ids = {
-        segment.segment_id
-        for segment in coordinator.data.state.segments.values()
-        if segment.segment_id is not None
-    }
-
     new_entities: list[WLEDNumber] = []
 
     # Process new segments, add them to Home Assistant
-    for segment_id in segment_ids - current_ids:
+    for segment_id in coordinator.segment_ids - current_ids:
         current_ids.add(segment_id)
         new_entities.extend(
             WLEDNumber(coordinator, segment_id, desc) for desc in NUMBERS

homeassistant/components/wled/select.py

@@ -208,16 +208,10 @@ def async_update_segments(
     async_add_entities: AddConfigEntryEntitiesCallback,
 ) -> None:
     """Update segments."""
-    segment_ids = {
-        segment.segment_id
-        for segment in coordinator.data.state.segments.values()
-        if segment.segment_id is not None
-    }
-
     new_entities: list[WLEDPaletteSelect] = []
 
     # Process new segments, add them to Home Assistant
-    for segment_id in segment_ids - current_ids:
+    for segment_id in coordinator.segment_ids - current_ids:
         current_ids.add(segment_id)
         new_entities.append(WLEDPaletteSelect(coordinator, segment_id))
 

homeassistant/components/wled/switch.py

@@ -245,16 +245,10 @@ def async_update_segments(
     async_add_entities: AddConfigEntryEntitiesCallback,
 ) -> None:
     """Update segments."""
-    segment_ids = {
-        segment.segment_id
-        for segment in coordinator.data.state.segments.values()
-        if segment.segment_id is not None
-    }
-
     new_entities: list[WLEDSegmentSwitch] = []
 
     # Process new segments, add them to Home Assistant
-    for segment_id in segment_ids - current_ids:
+    for segment_id in coordinator.segment_ids - current_ids:
         current_ids.add(segment_id)
         new_entities.extend(
             WLEDSegmentSwitch(

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"]