From 8af545a4e3f3d5cc4af3e3af7bbfd99e6319a20d Mon Sep 17 00:00:00 2001
From: Maciej Bieniek <bieniu@users.noreply.github.com>
Date: Sat, 1 Jan 2022 20:42:17 +0100
Subject: [PATCH] Add new buttons for Shelly Gas (#63099)

* Add self test/mute/unmute buttons

* Improve doctrings

* Improve test coverage

* Add supported property to ShellyButtonDescription

* Change icon
---
 homeassistant/components/shelly/button.py | 43 ++++++++++++++++++++---
 homeassistant/components/shelly/const.py  |  2 ++
 2 files changed, 40 insertions(+), 5 deletions(-)

diff --git a/homeassistant/components/shelly/button.py b/homeassistant/components/shelly/button.py
index 84c86f3fe91..f746d824984 100644
--- a/homeassistant/components/shelly/button.py
+++ b/homeassistant/components/shelly/button.py
@@ -18,7 +18,7 @@ from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.util import slugify
 
 from . import BlockDeviceWrapper, RpcDeviceWrapper
-from .const import BLOCK, DATA_CONFIG_ENTRY, DOMAIN, RPC
+from .const import BLOCK, DATA_CONFIG_ENTRY, DOMAIN, RPC, SHELLY_GAS_MODELS
 from .utils import get_block_device_name, get_device_entry_gen, get_rpc_device_name
 
 
@@ -33,6 +33,8 @@ class ShellyButtonDescriptionMixin:
 class ShellyButtonDescription(ButtonEntityDescription, ShellyButtonDescriptionMixin):
     """Class to describe a Button entity."""
 
+    supported: Callable = lambda _: True
+
 
 BUTTONS: Final = [
     ShellyButtonDescription(
@@ -57,6 +59,30 @@ BUTTONS: Final = [
         entity_category=EntityCategory.CONFIG,
         press_action=lambda wrapper: wrapper.device.trigger_reboot(),
     ),
+    ShellyButtonDescription(
+        key="self_test",
+        name="Self Test",
+        icon="mdi:progress-wrench",
+        entity_category=EntityCategory.DIAGNOSTIC,
+        press_action=lambda wrapper: wrapper.device.trigger_shelly_gas_self_test(),
+        supported=lambda wrapper: wrapper.device.model in SHELLY_GAS_MODELS,
+    ),
+    ShellyButtonDescription(
+        key="mute",
+        name="Mute",
+        icon="mdi:volume-mute",
+        entity_category=EntityCategory.CONFIG,
+        press_action=lambda wrapper: wrapper.device.trigger_shelly_gas_mute(),
+        supported=lambda wrapper: wrapper.device.model in SHELLY_GAS_MODELS,
+    ),
+    ShellyButtonDescription(
+        key="unmute",
+        name="Unmute",
+        icon="mdi:volume-high",
+        entity_category=EntityCategory.CONFIG,
+        press_action=lambda wrapper: wrapper.device.trigger_shelly_gas_unmute(),
+        supported=lambda wrapper: wrapper.device.model in SHELLY_GAS_MODELS,
+    ),
 ]
 
 
@@ -79,11 +105,18 @@ async def async_setup_entry(
             wrapper = cast(BlockDeviceWrapper, block_wrapper)
 
     if wrapper is not None:
-        async_add_entities([ShellyButton(wrapper, button) for button in BUTTONS])
+        entities = []
+
+        for button in BUTTONS:
+            if not button.supported(wrapper):
+                continue
+            entities.append(ShellyButton(wrapper, button))
+
+        async_add_entities(entities)
 
 
 class ShellyButton(ButtonEntity):
-    """Defines a Shelly OTA update base button."""
+    """Defines a Shelly base button."""
 
     entity_description: ShellyButtonDescription
 
@@ -92,7 +125,7 @@ class ShellyButton(ButtonEntity):
         wrapper: RpcDeviceWrapper | BlockDeviceWrapper,
         description: ShellyButtonDescription,
     ) -> None:
-        """Initialize Shelly OTA update button."""
+        """Initialize Shelly button."""
         self.entity_description = description
         self.wrapper = wrapper
 
@@ -108,5 +141,5 @@ class ShellyButton(ButtonEntity):
         )
 
     async def async_press(self) -> None:
-        """Triggers the OTA update service."""
+        """Triggers the Shelly button press service."""
         await self.entity_description.press_action(self.wrapper)
diff --git a/homeassistant/components/shelly/const.py b/homeassistant/components/shelly/const.py
index 8ef3ed5f1ac..68053ff0663 100644
--- a/homeassistant/components/shelly/const.py
+++ b/homeassistant/components/shelly/const.py
@@ -161,3 +161,5 @@ MAX_RPC_KEY_INSTANCES = 4
 
 # Time to wait before reloading entry upon device config change
 ENTRY_RELOAD_COOLDOWN = 60
+
+SHELLY_GAS_MODELS = ["SHGS-1"]