From e4f7ac62360c555cb392ce9e972373c34ded3202 Mon Sep 17 00:00:00 2001
From: Simone Chemelli <simone.chemelli@gmail.com>
Date: Sun, 13 Oct 2024 15:11:40 +0200
Subject: [PATCH] Add switch entity for Shelly scripts (#108171)

* introduce script switch only

* chore: add script switch test

* chore: apply review comments

* chore: fix tests

* chore: apply review comments
---
 homeassistant/components/shelly/switch.py | 46 ++++++++++++++++++
 tests/components/shelly/test_switch.py    | 59 +++++++++++++++++++++++
 2 files changed, 105 insertions(+)

diff --git a/homeassistant/components/shelly/switch.py b/homeassistant/components/shelly/switch.py
index 5ec223f53ad..134704cb0ff 100644
--- a/homeassistant/components/shelly/switch.py
+++ b/homeassistant/components/shelly/switch.py
@@ -66,6 +66,13 @@ RPC_VIRTUAL_SWITCH = RpcSwitchDescription(
     sub_key="value",
 )
 
+RPC_SCRIPT_SWITCH = RpcSwitchDescription(
+    key="script",
+    sub_key="running",
+    entity_registry_enabled_default=False,
+    entity_category=EntityCategory.CONFIG,
+)
+
 
 async def async_setup_entry(
     hass: HomeAssistant,
@@ -176,6 +183,14 @@ def async_setup_rpc_entry(
         RpcVirtualSwitch,
     )
 
+    async_setup_rpc_attribute_entities(
+        hass,
+        config_entry,
+        async_add_entities,
+        {"script": RPC_SCRIPT_SWITCH},
+        RpcScriptSwitch,
+    )
+
     # the user can remove virtual components from the device configuration, so we need
     # to remove orphaned entities
     virtual_switch_ids = get_virtual_component_ids(
@@ -190,6 +205,17 @@ def async_setup_rpc_entry(
         "boolean",
     )
 
+    # if the script is removed, from the device configuration, we need
+    # to remove orphaned entities
+    async_remove_orphaned_entities(
+        hass,
+        config_entry.entry_id,
+        coordinator.mac,
+        SWITCH_PLATFORM,
+        coordinator.device.status,
+        "script",
+    )
+
     if not switch_ids:
         return
 
@@ -317,3 +343,23 @@ class RpcVirtualSwitch(ShellyRpcAttributeEntity, SwitchEntity):
     async def async_turn_off(self, **kwargs: Any) -> None:
         """Turn off relay."""
         await self.call_rpc("Boolean.Set", {"id": self._id, "value": False})
+
+
+class RpcScriptSwitch(ShellyRpcAttributeEntity, SwitchEntity):
+    """Entity that controls a script component on RPC based Shelly devices."""
+
+    entity_description: RpcSwitchDescription
+    _attr_has_entity_name = True
+
+    @property
+    def is_on(self) -> bool:
+        """If switch is on."""
+        return bool(self.status["running"])
+
+    async def async_turn_on(self, **kwargs: Any) -> None:
+        """Turn on relay."""
+        await self.call_rpc("Script.Start", {"id": self._id})
+
+    async def async_turn_off(self, **kwargs: Any) -> None:
+        """Turn off relay."""
+        await self.call_rpc("Script.Stop", {"id": self._id})
diff --git a/tests/components/shelly/test_switch.py b/tests/components/shelly/test_switch.py
index c891d1d7b2d..5c7933afd7e 100644
--- a/tests/components/shelly/test_switch.py
+++ b/tests/components/shelly/test_switch.py
@@ -572,3 +572,62 @@ async def test_rpc_remove_virtual_switch_when_orphaned(
 
     entry = entity_registry.async_get(entity_id)
     assert not entry
+
+
+@pytest.mark.usefixtures("entity_registry_enabled_by_default")
+async def test_rpc_device_script_switch(
+    hass: HomeAssistant,
+    entity_registry: er.EntityRegistry,
+    mock_rpc_device: Mock,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """Test a script switch for RPC device."""
+    config = deepcopy(mock_rpc_device.config)
+    key = "script:1"
+    script_name = "aioshelly_ble_integration"
+    entity_id = f"switch.test_name_{script_name}"
+    config[key] = {
+        "id": 1,
+        "name": script_name,
+        "enable": False,
+    }
+    monkeypatch.setattr(mock_rpc_device, "config", config)
+
+    status = deepcopy(mock_rpc_device.status)
+    status[key] = {
+        "running": True,
+    }
+    monkeypatch.setattr(mock_rpc_device, "status", status)
+
+    await init_integration(hass, 3)
+
+    state = hass.states.get(entity_id)
+    assert state
+    assert state.state == STATE_ON
+    entry = entity_registry.async_get(entity_id)
+    assert entry
+    assert entry.unique_id == f"123456789ABC-{key}-script"
+
+    monkeypatch.setitem(mock_rpc_device.status[key], "running", False)
+    await hass.services.async_call(
+        SWITCH_DOMAIN,
+        SERVICE_TURN_OFF,
+        {ATTR_ENTITY_ID: entity_id},
+        blocking=True,
+    )
+    mock_rpc_device.mock_update()
+    state = hass.states.get(entity_id)
+    assert state
+    assert state.state == STATE_OFF
+
+    monkeypatch.setitem(mock_rpc_device.status[key], "running", True)
+    await hass.services.async_call(
+        SWITCH_DOMAIN,
+        SERVICE_TURN_ON,
+        {ATTR_ENTITY_ID: entity_id},
+        blocking=True,
+    )
+    mock_rpc_device.mock_update()
+    state = hass.states.get(entity_id)
+    assert state
+    assert state.state == STATE_ON