diff --git a/homeassistant/components/wled/coordinator.py b/homeassistant/components/wled/coordinator.py index 2ab034f421c3ff..d94fadad50ffff 100644 --- a/homeassistant/components/wled/coordinator.py +++ b/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.""" diff --git a/homeassistant/components/wled/light.py b/homeassistant/components/wled/light.py index c82884ebace6fe..701f765ebaa95e 100644 --- a/homeassistant/components/wled/light.py +++ b/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 diff --git a/homeassistant/components/wled/number.py b/homeassistant/components/wled/number.py index 3964a0c7dd0c7e..bfc5dce12ed10f 100644 --- a/homeassistant/components/wled/number.py +++ b/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 diff --git a/homeassistant/components/wled/select.py b/homeassistant/components/wled/select.py index 70eb8e5a901414..a2d21d8e3e704d 100644 --- a/homeassistant/components/wled/select.py +++ b/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)) diff --git a/homeassistant/components/wled/switch.py b/homeassistant/components/wled/switch.py index 6b5b76147801a4..155732dad5f26f 100644 --- a/homeassistant/components/wled/switch.py +++ b/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( diff --git a/tests/components/wled/snapshots/test_light.ambr b/tests/components/wled/snapshots/test_light.ambr index 1d49979e407d08..33f13c706977ef 100644 --- a/tests/components/wled/snapshots/test_light.ambr +++ b/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([ , ]), @@ -958,6 +962,10 @@ 'brightness': 128, 'color_mode': , 'friendly_name': 'WLED RGB Light Main', + 'group_entities': list([ + 'light.wled_rgb_light', + 'light.wled_rgb_light_segment_1', + ]), 'supported_color_modes': list([ , ]), diff --git a/tests/components/wled/test_light.py b/tests/components/wled/test_light.py index 412720ae89ea8e..5f64a8433eaa32 100644 --- a/tests/components/wled/test_light.py +++ b/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"]