Switch cleanup for Shelly (part 2) (#138922)

* Switch cleanup for Shelly (part 2)

* apply review comment

* Update tests/components/shelly/test_climate.py

Co-authored-by: Maciej Bieniek <bieniu@users.noreply.github.com>

* apply review comments

---------

Co-authored-by: Maciej Bieniek <bieniu@users.noreply.github.com>

homeassistant/components/shelly/switch.py

@@ -7,8 +7,9 @@ from dataclasses import dataclass
 from typing import Any, cast
 
 from aioshelly.block_device import Block
-from aioshelly.const import MODEL_2, MODEL_25, MODEL_WALL_DISPLAY, RPC_GENERATIONS
+from aioshelly.const import MODEL_2, MODEL_25, RPC_GENERATIONS
 
+from homeassistant.components.climate import DOMAIN as CLIMATE_PLATFORM
 from homeassistant.components.switch import (
     DOMAIN as SWITCH_PLATFORM,
     SwitchEntity,
@@ -27,7 +28,6 @@ from .entity import (
     RpcEntityDescription,
     ShellyBlockEntity,
     ShellyRpcAttributeEntity,
-    ShellyRpcEntity,
     ShellySleepingBlockAttributeEntity,
     async_setup_entry_attribute_entities,
     async_setup_entry_rpc,
@@ -36,12 +36,9 @@ from .utils import (
     async_remove_orphaned_entities,
     async_remove_shelly_entity,
     get_device_entry_gen,
-    get_rpc_key_ids,
     get_virtual_component_ids,
     is_block_channel_type_light,
-    is_rpc_channel_type_light,
+    is_rpc_exclude_from_relay,
-    is_rpc_thermostat_internal_actuator,
-    is_rpc_thermostat_mode,
 )
 
 
@@ -67,6 +64,18 @@ class RpcSwitchDescription(RpcEntityDescription, SwitchEntityDescription):
     method_params_fn: Callable[[int | None, bool], dict]
 
 
+RPC_RELAY_SWITCHES = {
+    "switch": RpcSwitchDescription(
+        key="switch",
+        sub_key="output",
+        removal_condition=is_rpc_exclude_from_relay,
+        is_on=lambda status: bool(status["output"]),
+        method_on="Switch.Set",
+        method_off="Switch.Set",
+        method_params_fn=lambda id, value: {"id": id, "on": value},
+    ),
+}
+
 RPC_SWITCHES = {
     "boolean": RpcSwitchDescription(
         key="boolean",
@@ -162,32 +171,10 @@ def async_setup_rpc_entry(
     """Set up entities for RPC device."""
     coordinator = config_entry.runtime_data.rpc
     assert coordinator
-    switch_key_ids = get_rpc_key_ids(coordinator.device.status, "switch")
 
-    switch_ids = []
+    async_setup_entry_rpc(
-    for id_ in switch_key_ids:
+        hass, config_entry, async_add_entities, RPC_RELAY_SWITCHES, RpcRelaySwitch
-        if is_rpc_channel_type_light(coordinator.device.config, id_):
+    )
-            continue
-
-        if coordinator.model == MODEL_WALL_DISPLAY:
-            # There are three configuration scenarios for WallDisplay:
-            # - relay mode (no thermostat)
-            # - thermostat mode using the internal relay as an actuator
-            # - thermostat mode using an external (from another device) relay as
-            #   an actuator
-            if not is_rpc_thermostat_mode(id_, coordinator.device.status):
-                # The device is not in thermostat mode, we need to remove a climate
-                # entity
-                unique_id = f"{coordinator.mac}-thermostat:{id_}"
-                async_remove_shelly_entity(hass, "climate", unique_id)
-            elif is_rpc_thermostat_internal_actuator(coordinator.device.status):
-                # The internal relay is an actuator, skip this ID so as not to create
-                # a switch entity
-                continue
-
-        switch_ids.append(id_)
-        unique_id = f"{coordinator.mac}-switch:{id_}"
-        async_remove_shelly_entity(hass, "light", unique_id)
 
     async_setup_entry_rpc(
         hass, config_entry, async_add_entities, RPC_SWITCHES, RpcSwitch
@@ -218,10 +205,16 @@ def async_setup_rpc_entry(
         "script",
     )
 
-    if not switch_ids:
+    # if the climate is removed, from the device configuration, we need
-        return
+    # to remove orphaned entities
-
+    async_remove_orphaned_entities(
-    async_add_entities(RpcRelaySwitch(coordinator, id_) for id_ in switch_ids)
+        hass,
+        config_entry.entry_id,
+        coordinator.mac,
+        CLIMATE_PLATFORM,
+        coordinator.device.status,
+        "thermostat",
+    )
 
 
 class BlockSleepingMotionSwitch(
@@ -305,28 +298,6 @@ class BlockRelaySwitch(ShellyBlockEntity, SwitchEntity):
         super()._update_callback()
 
 
-class RpcRelaySwitch(ShellyRpcEntity, SwitchEntity):
-    """Entity that controls a relay on RPC based Shelly devices."""
-
-    def __init__(self, coordinator: ShellyRpcCoordinator, id_: int) -> None:
-        """Initialize relay switch."""
-        super().__init__(coordinator, f"switch:{id_}")
-        self._id = id_
-
-    @property
-    def is_on(self) -> bool:
-        """If switch is on."""
-        return bool(self.status["output"])
-
-    async def async_turn_on(self, **kwargs: Any) -> None:
-        """Turn on relay."""
-        await self.call_rpc("Switch.Set", {"id": self._id, "on": True})
-
-    async def async_turn_off(self, **kwargs: Any) -> None:
-        """Turn off relay."""
-        await self.call_rpc("Switch.Set", {"id": self._id, "on": False})
-
-
 class RpcSwitch(ShellyRpcAttributeEntity, SwitchEntity):
     """Entity that controls a switch on RPC based Shelly devices."""
 
@@ -351,3 +322,21 @@ class RpcSwitch(ShellyRpcAttributeEntity, SwitchEntity):
             self.entity_description.method_off,
             self.entity_description.method_params_fn(self._id, False),
         )
+
+
+class RpcRelaySwitch(RpcSwitch):
+    """Entity that controls a switch on RPC based Shelly devices."""
+
+    # False to avoid double naming as True is inerithed from base class
+    _attr_has_entity_name = False
+
+    def __init__(
+        self,
+        coordinator: ShellyRpcCoordinator,
+        key: str,
+        attribute: str,
+        description: RpcEntityDescription,
+    ) -> None:
+        """Initialize the switch."""
+        super().__init__(coordinator, key, attribute, description)
+        self._attr_unique_id: str = f"{coordinator.mac}-{key}"

homeassistant/components/shelly/utils.py

@@ -627,3 +627,14 @@ async def get_rpc_script_event_types(device: RpcDevice, id: int) -> list[str]:
     code_response = await device.script_getcode(id)
     matches = SHELLY_EMIT_EVENT_PATTERN.finditer(code_response["data"])
     return sorted([*{str(event_type.group(1)) for event_type in matches}])
+
+
+def is_rpc_exclude_from_relay(
+    settings: dict[str, Any], status: dict[str, Any], channel: str
+) -> bool:
+    """Return true if rpc channel should be excludeed from switch platform."""
+    ch = int(channel.split(":")[1])
+    if is_rpc_thermostat_internal_actuator(status):
+        return True
+
+    return is_rpc_channel_type_light(settings, ch)

tests/components/shelly/conftest.py

@@ -101,6 +101,7 @@ MOCK_BLOCKS = [
             "overpower": 0,
             "power": 53.4,
             "energy": 1234567.89,
+            "output": True,
         },
         channel="0",
         type="relay",
@@ -207,7 +208,7 @@ MOCK_CONFIG = {
     },
     "sys": {
         "ui_data": {},
-        "device": {"name": "Test name"},
+        "device": {"name": "Test name", "mac": MOCK_MAC},
     },
     "wifi": {"sta": {"enable": True}, "sta1": {"enable": False}},
     "ws": {"enable": False, "server": None},
@@ -312,7 +313,11 @@ MOCK_STATUS_COAP = {
 
 
 MOCK_STATUS_RPC = {
-    "switch:0": {"output": True},
+    "switch:0": {
+        "id": 0,
+        "output": True,
+        "apower": 85.3,
+    },
     "input:0": {"id": 0, "state": None},
     "input:1": {"id": 1, "percent": 89, "xpercent": 8.9},
     "input:2": {

tests/components/shelly/test_climate.py

@@ -751,6 +751,7 @@ async def test_wall_display_thermostat_mode_external_actuator(
 
     new_status = deepcopy(mock_rpc_device.status)
     new_status["sys"]["relay_in_thermostat"] = False
+    new_status.pop("cover:0")
     monkeypatch.setattr(mock_rpc_device, "status", new_status)
 
     await init_integration(hass, 2, model=MODEL_WALL_DISPLAY)

tests/components/shelly/test_coordinator.py

@@ -386,6 +386,8 @@ async def test_rpc_reload_on_cfg_change(
     monkeypatch: pytest.MonkeyPatch,
 ) -> None:
     """Test RPC reload on config change."""
+    monkeypatch.delitem(mock_rpc_device.status, "cover:0")
+    monkeypatch.setitem(mock_rpc_device.status["sys"], "relay_in_thermostat", False)
     await init_integration(hass, 2)
 
     # Generate config change from switch to light
@@ -710,6 +712,8 @@ async def test_rpc_reconnect_error(
     exc: Exception,
 ) -> None:
     """Test RPC reconnect error."""
+    monkeypatch.delitem(mock_rpc_device.status, "cover:0")
+    monkeypatch.setitem(mock_rpc_device.status["sys"], "relay_in_thermostat", False)
     await init_integration(hass, 2)
 
     assert get_entity_state(hass, "switch.test_switch_0") == STATE_ON
@@ -729,9 +733,12 @@ async def test_rpc_error_running_connected_events(
     hass: HomeAssistant,
     freezer: FrozenDateTimeFactory,
     mock_rpc_device: Mock,
+    monkeypatch: pytest.MonkeyPatch,
     caplog: pytest.LogCaptureFixture,
 ) -> None:
     """Test RPC error while running connected events."""
+    monkeypatch.delitem(mock_rpc_device.status, "cover:0")
+    monkeypatch.setitem(mock_rpc_device.status["sys"], "relay_in_thermostat", False)
     with patch(
         "homeassistant.components.shelly.coordinator.async_ensure_ble_enabled",
         side_effect=DeviceConnectionError,

tests/components/shelly/test_init.py

@@ -366,8 +366,11 @@ async def test_entry_unload(
     entity_id: str,
     mock_block_device: Mock,
     mock_rpc_device: Mock,
+    monkeypatch: pytest.MonkeyPatch,
 ) -> None:
     """Test entry unload."""
+    monkeypatch.delitem(mock_rpc_device.status, "cover:0")
+    monkeypatch.setitem(mock_rpc_device.status["sys"], "relay_in_thermostat", False)
     entry = await init_integration(hass, gen)
 
     assert entry.state is ConfigEntryState.LOADED
@@ -410,6 +413,9 @@ async def test_entry_unload_not_connected(
     hass: HomeAssistant, mock_rpc_device: Mock, monkeypatch: pytest.MonkeyPatch
 ) -> None:
     """Test entry unload when not connected."""
+    monkeypatch.delitem(mock_rpc_device.status, "cover:0")
+    monkeypatch.setitem(mock_rpc_device.status["sys"], "relay_in_thermostat", False)
+
     with patch(
         "homeassistant.components.shelly.coordinator.async_stop_scanner"
     ) as mock_stop_scanner:
@@ -435,6 +441,9 @@ async def test_entry_unload_not_connected_but_we_think_we_are(
     hass: HomeAssistant, mock_rpc_device: Mock, monkeypatch: pytest.MonkeyPatch
 ) -> None:
     """Test entry unload when not connected but we think we are still connected."""
+    monkeypatch.delitem(mock_rpc_device.status, "cover:0")
+    monkeypatch.setitem(mock_rpc_device.status["sys"], "relay_in_thermostat", False)
+
     with patch(
         "homeassistant.components.shelly.coordinator.async_stop_scanner",
         side_effect=DeviceConnectionError,

tests/components/shelly/test_switch.py

@@ -288,6 +288,8 @@ async def test_rpc_device_services(
     hass: HomeAssistant, mock_rpc_device: Mock, monkeypatch: pytest.MonkeyPatch
 ) -> None:
     """Test RPC device turn on/off services."""
+    monkeypatch.delitem(mock_rpc_device.status, "cover:0")
+    monkeypatch.setitem(mock_rpc_device.status["sys"], "relay_in_thermostat", False)
     await init_integration(hass, 2)
 
     await hass.services.async_call(
@@ -310,9 +312,14 @@ async def test_rpc_device_services(
 
 
 async def test_rpc_device_unique_ids(
-    hass: HomeAssistant, mock_rpc_device: Mock, entity_registry: EntityRegistry
+    hass: HomeAssistant,
+    mock_rpc_device: Mock,
+    monkeypatch: pytest.MonkeyPatch,
+    entity_registry: EntityRegistry,
 ) -> None:
     """Test RPC device unique_ids."""
+    monkeypatch.delitem(mock_rpc_device.status, "cover:0")
+    monkeypatch.setitem(mock_rpc_device.status["sys"], "relay_in_thermostat", False)
     await init_integration(hass, 2)
 
     entry = entity_registry.async_get("switch.test_switch_0")
@@ -340,6 +347,8 @@ async def test_rpc_set_state_errors(
 ) -> None:
     """Test RPC device set state connection/call errors."""
     monkeypatch.setattr(mock_rpc_device, "call_rpc", AsyncMock(side_effect=exc))
+    monkeypatch.delitem(mock_rpc_device.status, "cover:0")
+    monkeypatch.setitem(mock_rpc_device.status["sys"], "relay_in_thermostat", False)
     await init_integration(hass, 2)
 
     with pytest.raises(HomeAssistantError):
@@ -360,6 +369,8 @@ async def test_rpc_auth_error(
         "call_rpc",
         AsyncMock(side_effect=InvalidAuthError),
     )
+    monkeypatch.delitem(mock_rpc_device.status, "cover:0")
+    monkeypatch.setitem(mock_rpc_device.status["sys"], "relay_in_thermostat", False)
     entry = await init_integration(hass, 2)
 
     assert entry.state is ConfigEntryState.LOADED
@@ -409,15 +420,22 @@ async def test_wall_display_relay_mode(
     monkeypatch: pytest.MonkeyPatch,
 ) -> None:
     """Test Wall Display in relay mode."""
-    climate_entity_id = "climate.test_name"
+    climate_entity_id = "climate.test_name_thermostat_0"
     switch_entity_id = "switch.test_switch_0"
 
+    config_entry = await init_integration(hass, 2, model=MODEL_WALL_DISPLAY)
+
+    assert hass.states.get(climate_entity_id) is not None
+    assert len(hass.states.async_entity_ids(CLIMATE_DOMAIN)) == 1
+
     new_status = deepcopy(mock_rpc_device.status)
     new_status["sys"]["relay_in_thermostat"] = False
     new_status.pop("thermostat:0")
+    new_status.pop("cover:0")
     monkeypatch.setattr(mock_rpc_device, "status", new_status)
 
-    await init_integration(hass, 2, model=MODEL_WALL_DISPLAY)
+    await hass.config_entries.async_reload(config_entry.entry_id)
+    await hass.async_block_till_done()
 
     # the climate entity should be removed
     assert hass.states.get(climate_entity_id) is None