From 3e99afdd540b6ddc6941e0bdf5b393f9960fce78 Mon Sep 17 00:00:00 2001 From: Shay Levy Date: Sat, 30 Mar 2024 18:48:57 +0300 Subject: [PATCH] Cleanup Shelly RGBW light entities (#114410) --- homeassistant/components/shelly/const.py | 2 + homeassistant/components/shelly/light.py | 17 +++++ homeassistant/components/shelly/utils.py | 12 ++++ tests/components/shelly/__init__.py | 12 ++++ tests/components/shelly/conftest.py | 6 ++ tests/components/shelly/test_light.py | 88 +++++++++++++++++++++++- 6 files changed, 134 insertions(+), 3 deletions(-) diff --git a/homeassistant/components/shelly/const.py b/homeassistant/components/shelly/const.py index 3580bcf9b38..2ac0416bb6c 100644 --- a/homeassistant/components/shelly/const.py +++ b/homeassistant/components/shelly/const.py @@ -234,3 +234,5 @@ DEVICES_WITHOUT_FIRMWARE_CHANGELOG = ( ) CONF_GEN = "gen" + +SHELLY_PLUS_RGBW_CHANNELS = 4 diff --git a/homeassistant/components/shelly/light.py b/homeassistant/components/shelly/light.py index 6c28023a5e3..d0590fc7c20 100644 --- a/homeassistant/components/shelly/light.py +++ b/homeassistant/components/shelly/light.py @@ -14,6 +14,7 @@ from homeassistant.components.light import ( ATTR_RGB_COLOR, ATTR_RGBW_COLOR, ATTR_TRANSITION, + DOMAIN as LIGHT_DOMAIN, ColorMode, LightEntity, LightEntityFeature, @@ -34,12 +35,14 @@ from .const import ( RGBW_MODELS, RPC_MIN_TRANSITION_TIME_SEC, SHBLB_1_RGB_EFFECTS, + SHELLY_PLUS_RGBW_CHANNELS, STANDARD_RGB_EFFECTS, ) from .coordinator import ShellyBlockCoordinator, ShellyRpcCoordinator, get_entry_data from .entity import ShellyBlockEntity, ShellyRpcEntity from .utils import ( async_remove_shelly_entity, + async_remove_shelly_rpc_entities, brightness_to_percentage, get_device_entry_gen, get_rpc_key_ids, @@ -118,14 +121,28 @@ def async_setup_rpc_entry( return if light_key_ids := get_rpc_key_ids(coordinator.device.status, "light"): + # Light mode remove RGB & RGBW entities, add light entities + async_remove_shelly_rpc_entities( + hass, LIGHT_DOMAIN, coordinator.mac, ["rgb:0", "rgbw:0"] + ) async_add_entities(RpcShellyLight(coordinator, id_) for id_ in light_key_ids) return + light_keys = [f"light:{i}" for i in range(SHELLY_PLUS_RGBW_CHANNELS)] + if rgb_key_ids := get_rpc_key_ids(coordinator.device.status, "rgb"): + # RGB mode remove light & RGBW entities, add RGB entity + async_remove_shelly_rpc_entities( + hass, LIGHT_DOMAIN, coordinator.mac, [*light_keys, "rgbw:0"] + ) async_add_entities(RpcShellyRgbLight(coordinator, id_) for id_ in rgb_key_ids) return if rgbw_key_ids := get_rpc_key_ids(coordinator.device.status, "rgbw"): + # RGBW mode remove light & RGB entities, add RGBW entity + async_remove_shelly_rpc_entities( + hass, LIGHT_DOMAIN, coordinator.mac, [*light_keys, "rgb:0"] + ) async_add_entities(RpcShellyRgbwLight(coordinator, id_) for id_ in rgbw_key_ids) diff --git a/homeassistant/components/shelly/utils.py b/homeassistant/components/shelly/utils.py index d26e3dc11f3..ce98e0d5c12 100644 --- a/homeassistant/components/shelly/utils.py +++ b/homeassistant/components/shelly/utils.py @@ -488,3 +488,15 @@ async def async_shutdown_device(device: BlockDevice | RpcDevice) -> None: await device.shutdown() if isinstance(device, BlockDevice): device.shutdown() + + +@callback +def async_remove_shelly_rpc_entities( + hass: HomeAssistant, domain: str, mac: str, keys: list[str] +) -> None: + """Remove RPC based Shelly entity.""" + entity_reg = er_async_get(hass) + for key in keys: + if entity_id := entity_reg.async_get_entity_id(domain, DOMAIN, f"{mac}-{key}"): + LOGGER.debug("Removing entity: %s", entity_id) + entity_reg.async_remove(entity_id) diff --git a/tests/components/shelly/__init__.py b/tests/components/shelly/__init__.py index 2dc9012d863..348b1115a6f 100644 --- a/tests/components/shelly/__init__.py +++ b/tests/components/shelly/__init__.py @@ -126,6 +126,18 @@ def register_entity( return f"{domain}.{object_id}" +def get_entity( + hass: HomeAssistant, + domain: str, + unique_id: str, +) -> str | None: + """Get Shelly entity.""" + entity_registry = async_get(hass) + return entity_registry.async_get_entity_id( + domain, DOMAIN, f"{MOCK_MAC}-{unique_id}" + ) + + def get_entity_state(hass: HomeAssistant, entity_id: str) -> str: """Return entity state.""" entity = hass.states.get(entity_id) diff --git a/tests/components/shelly/conftest.py b/tests/components/shelly/conftest.py index 9a73252ca6c..3cd27101f76 100644 --- a/tests/components/shelly/conftest.py +++ b/tests/components/shelly/conftest.py @@ -169,6 +169,9 @@ MOCK_CONFIG = { "input:1": {"id": 1, "type": "analog", "enable": True}, "input:2": {"id": 2, "name": "Gas", "type": "count", "enable": True}, "light:0": {"name": "test light_0"}, + "light:1": {"name": "test light_1"}, + "light:2": {"name": "test light_2"}, + "light:3": {"name": "test light_3"}, "rgb:0": {"name": "test rgb_0"}, "rgbw:0": {"name": "test rgbw_0"}, "switch:0": {"name": "test switch_0"}, @@ -225,6 +228,9 @@ MOCK_STATUS_RPC = { "input:1": {"id": 1, "percent": 89, "xpercent": 8.9}, "input:2": {"id": 2, "counts": {"total": 56174, "xtotal": 561.74}}, "light:0": {"output": True, "brightness": 53.0}, + "light:1": {"output": True, "brightness": 53.0}, + "light:2": {"output": True, "brightness": 53.0}, + "light:3": {"output": True, "brightness": 53.0}, "rgb:0": {"output": True, "brightness": 53.0, "rgb": [45, 55, 65]}, "rgbw:0": {"output": True, "brightness": 53.0, "rgb": [21, 22, 23], "white": 120}, "cloud": {"connected": False}, diff --git a/tests/components/shelly/test_light.py b/tests/components/shelly/test_light.py index cca318c364d..2c464a8c39c 100644 --- a/tests/components/shelly/test_light.py +++ b/tests/components/shelly/test_light.py @@ -29,6 +29,7 @@ from homeassistant.components.light import ( ColorMode, LightEntityFeature, ) +from homeassistant.components.shelly.const import SHELLY_PLUS_RGBW_CHANNELS from homeassistant.const import ( ATTR_ENTITY_ID, ATTR_SUPPORTED_FEATURES, @@ -38,7 +39,7 @@ from homeassistant.const import ( from homeassistant.core import HomeAssistant from homeassistant.helpers.entity_registry import EntityRegistry -from . import init_integration, mutate_rpc_device_status +from . import get_entity, init_integration, mutate_rpc_device_status, register_entity from .conftest import mock_white_light_set_state RELAY_BLOCK_ID = 0 @@ -587,7 +588,8 @@ async def test_rpc_device_rgb_profile( monkeypatch: pytest.MonkeyPatch, ) -> None: """Test RPC device in RGB profile.""" - monkeypatch.delitem(mock_rpc_device.status, "light:0") + for i in range(SHELLY_PLUS_RGBW_CHANNELS): + monkeypatch.delitem(mock_rpc_device.status, f"light:{i}") monkeypatch.delitem(mock_rpc_device.status, "rgbw:0") entity_id = "light.test_rgb_0" await init_integration(hass, 2) @@ -633,7 +635,8 @@ async def test_rpc_device_rgbw_profile( monkeypatch: pytest.MonkeyPatch, ) -> None: """Test RPC device in RGBW profile.""" - monkeypatch.delitem(mock_rpc_device.status, "light:0") + for i in range(SHELLY_PLUS_RGBW_CHANNELS): + monkeypatch.delitem(mock_rpc_device.status, f"light:{i}") monkeypatch.delitem(mock_rpc_device.status, "rgb:0") entity_id = "light.test_rgbw_0" await init_integration(hass, 2) @@ -673,3 +676,82 @@ async def test_rpc_device_rgbw_profile( entry = entity_registry.async_get(entity_id) assert entry assert entry.unique_id == "123456789ABC-rgbw:0" + + +async def test_rpc_rgbw_device_light_mode_remove_others( + hass: HomeAssistant, + mock_rpc_device: Mock, + entity_registry: EntityRegistry, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Test Shelly RPC RGBW device in light mode removes RGB/RGBW entities.""" + # register lights + monkeypatch.delitem(mock_rpc_device.status, "rgb:0") + monkeypatch.delitem(mock_rpc_device.status, "rgbw:0") + register_entity(hass, LIGHT_DOMAIN, "test_rgb_0", "rgb:0") + register_entity(hass, LIGHT_DOMAIN, "test_rgbw_0", "rgbw:0") + + # verify RGB & RGBW entities created + assert get_entity(hass, LIGHT_DOMAIN, "rgb:0") is not None + assert get_entity(hass, LIGHT_DOMAIN, "rgbw:0") is not None + + # init to remove RGB & RGBW + await init_integration(hass, 2) + + # verify we have 4 lights + for i in range(SHELLY_PLUS_RGBW_CHANNELS): + entity_id = f"light.test_light_{i}" + assert hass.states.get(entity_id).state == STATE_ON + entry = entity_registry.async_get(entity_id) + assert entry + assert entry.unique_id == f"123456789ABC-light:{i}" + + # verify RGB & RGBW entities removed + assert get_entity(hass, LIGHT_DOMAIN, "rgb:0") is None + assert get_entity(hass, LIGHT_DOMAIN, "rgbw:0") is None + + +@pytest.mark.parametrize( + ("active_mode", "removed_mode"), + [ + ("rgb", "rgbw"), + ("rgbw", "rgb"), + ], +) +async def test_rpc_rgbw_device_rgb_w_modes_remove_others( + hass: HomeAssistant, + mock_rpc_device: Mock, + entity_registry: EntityRegistry, + monkeypatch: pytest.MonkeyPatch, + active_mode: str, + removed_mode: str, +) -> None: + """Test Shelly RPC RGBW device in RGB/W modes other lights.""" + removed_key = f"{removed_mode}:0" + + # register lights + for i in range(SHELLY_PLUS_RGBW_CHANNELS): + monkeypatch.delitem(mock_rpc_device.status, f"light:{i}") + entity_id = f"light.test_light_{i}" + register_entity(hass, LIGHT_DOMAIN, entity_id, f"light:{i}") + monkeypatch.delitem(mock_rpc_device.status, f"{removed_mode}:0") + register_entity(hass, LIGHT_DOMAIN, f"test_{removed_key}", removed_key) + + # verify lights entities created + for i in range(SHELLY_PLUS_RGBW_CHANNELS): + assert get_entity(hass, LIGHT_DOMAIN, f"light:{i}") is not None + assert get_entity(hass, LIGHT_DOMAIN, removed_key) is not None + + await init_integration(hass, 2) + + # verify we have RGB/w light + entity_id = f"light.test_{active_mode}_0" + assert hass.states.get(entity_id).state == STATE_ON + entry = entity_registry.async_get(entity_id) + assert entry + assert entry.unique_id == f"123456789ABC-{active_mode}:0" + + # verify light & RGB/W entities removed + for i in range(SHELLY_PLUS_RGBW_CHANNELS): + assert get_entity(hass, LIGHT_DOMAIN, f"light:{i}") is None + assert get_entity(hass, LIGHT_DOMAIN, removed_key) is None