From 542f637ac4310d3322fa12151c98841486dfb9e2 Mon Sep 17 00:00:00 2001
From: Shay Levy <levyshay1@gmail.com>
Date: Mon, 20 Sep 2021 21:12:18 +0300
Subject: [PATCH] Improve Shelly light application/consumption type handling
 (#56461)

---
 homeassistant/components/shelly/light.py  | 18 +++++++++++-------
 homeassistant/components/shelly/switch.py | 21 +++++++++++----------
 homeassistant/components/shelly/utils.py  | 12 ++++++++++++
 3 files changed, 34 insertions(+), 17 deletions(-)

diff --git a/homeassistant/components/shelly/light.py b/homeassistant/components/shelly/light.py
index 6a1035816a5..cd034c1e7e5 100644
--- a/homeassistant/components/shelly/light.py
+++ b/homeassistant/components/shelly/light.py
@@ -51,7 +51,13 @@ from .const import (
     STANDARD_RGB_EFFECTS,
 )
 from .entity import ShellyBlockEntity, ShellyRpcEntity
-from .utils import async_remove_shelly_entity, get_device_entry_gen, get_rpc_key_ids
+from .utils import (
+    async_remove_shelly_entity,
+    get_device_entry_gen,
+    get_rpc_key_ids,
+    is_block_channel_type_light,
+    is_rpc_channel_type_light,
+)
 
 _LOGGER: Final = logging.getLogger(__name__)
 
@@ -82,10 +88,9 @@ async def async_setup_block_entry(
         if block.type == "light":
             blocks.append(block)
         elif block.type == "relay":
-            app_type = wrapper.device.settings["relays"][int(block.channel)].get(
-                "appliance_type"
-            )
-            if not app_type or app_type.lower() != "light":
+            if not is_block_channel_type_light(
+                wrapper.device.settings, int(block.channel)
+            ):
                 continue
 
             blocks.append(block)
@@ -110,8 +115,7 @@ async def async_setup_rpc_entry(
 
     switch_ids = []
     for id_ in switch_key_ids:
-        con_types = wrapper.device.config["sys"]["ui_data"].get("consumption_types")
-        if con_types is None or con_types[id_] != "lights":
+        if not is_rpc_channel_type_light(wrapper.device.config, id_):
             continue
 
         switch_ids.append(id_)
diff --git a/homeassistant/components/shelly/switch.py b/homeassistant/components/shelly/switch.py
index d6e8fa11798..0291258b511 100644
--- a/homeassistant/components/shelly/switch.py
+++ b/homeassistant/components/shelly/switch.py
@@ -13,7 +13,13 @@ from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from . import BlockDeviceWrapper, RpcDeviceWrapper
 from .const import BLOCK, DATA_CONFIG_ENTRY, DOMAIN, RPC
 from .entity import ShellyBlockEntity, ShellyRpcEntity
-from .utils import async_remove_shelly_entity, get_device_entry_gen, get_rpc_key_ids
+from .utils import (
+    async_remove_shelly_entity,
+    get_device_entry_gen,
+    get_rpc_key_ids,
+    is_block_channel_type_light,
+    is_rpc_channel_type_light,
+)
 
 
 async def async_setup_entry(
@@ -46,13 +52,9 @@ async def async_setup_block_entry(
     relay_blocks = []
     assert wrapper.device.blocks
     for block in wrapper.device.blocks:
-        if block.type != "relay":
-            continue
-
-        app_type = wrapper.device.settings["relays"][int(block.channel)].get(
-            "appliance_type"
-        )
-        if app_type and app_type.lower() == "light":
+        if block.type != "relay" or is_block_channel_type_light(
+            wrapper.device.settings, int(block.channel)
+        ):
             continue
 
         relay_blocks.append(block)
@@ -76,8 +78,7 @@ async def async_setup_rpc_entry(
 
     switch_ids = []
     for id_ in switch_key_ids:
-        con_types = wrapper.device.config["sys"]["ui_data"].get("consumption_types")
-        if con_types is not None and con_types[id_] == "lights":
+        if is_rpc_channel_type_light(wrapper.device.config, id_):
             continue
 
         switch_ids.append(id_)
diff --git a/homeassistant/components/shelly/utils.py b/homeassistant/components/shelly/utils.py
index 13b34ef5aea..4d3655829a7 100644
--- a/homeassistant/components/shelly/utils.py
+++ b/homeassistant/components/shelly/utils.py
@@ -302,3 +302,15 @@ def get_rpc_key_ids(keys_dict: dict[str, Any], key: str) -> list[int]:
 def is_rpc_momentary_input(config: dict[str, Any], key: str) -> bool:
     """Return true if rpc input button settings is set to a momentary type."""
     return cast(bool, config[key]["type"] == "button")
+
+
+def is_block_channel_type_light(settings: dict[str, Any], channel: int) -> bool:
+    """Return true if block channel appliance type is set to light."""
+    app_type = settings["relays"][channel].get("appliance_type")
+    return app_type is not None and app_type.lower().startswith("light")
+
+
+def is_rpc_channel_type_light(config: dict[str, Any], channel: int) -> bool:
+    """Return true if rpc channel consumption type is set to light."""
+    con_types = config["sys"]["ui_data"].get("consumption_types")
+    return con_types is not None and con_types[channel].lower().startswith("light")