From 09aaf45f0a879c8ecdad8c0fa98398e59900c71b Mon Sep 17 00:00:00 2001
From: Kevin Addeman <kevin.addeman@gmail.com>
Date: Thu, 18 Aug 2022 22:23:20 -0400
Subject: [PATCH] Fix lutron caseta Sunnata Keypad support (#75324)

Co-authored-by: J. Nick Koston <nick@koston.org>
---
 .../components/lutron_caseta/__init__.py      |  3 +-
 .../lutron_caseta/device_trigger.py           | 31 +++++++++++++++----
 .../lutron_caseta/test_device_trigger.py      |  4 +--
 3 files changed, 29 insertions(+), 9 deletions(-)

diff --git a/homeassistant/components/lutron_caseta/__init__.py b/homeassistant/components/lutron_caseta/__init__.py
index d235f294b3a..5653504c98a 100644
--- a/homeassistant/components/lutron_caseta/__init__.py
+++ b/homeassistant/components/lutron_caseta/__init__.py
@@ -46,6 +46,7 @@ from .const import (
 from .device_trigger import (
     DEVICE_TYPE_SUBTYPE_MAP_TO_LIP,
     LEAP_TO_DEVICE_TYPE_SUBTYPE_MAP,
+    _lutron_model_to_device_type,
 )
 from .models import LutronCasetaData
 from .util import serial_to_unique_id
@@ -281,7 +282,7 @@ def _async_subscribe_pico_remote_events(
         else:
             action = ACTION_RELEASE
 
-        type_ = device["type"]
+        type_ = _lutron_model_to_device_type(device["model"], device["type"])
         area, name = _area_and_name_from_name(device["name"])
         leap_button_number = device["button_number"]
         lip_button_number = async_get_lip_button(type_, leap_button_number)
diff --git a/homeassistant/components/lutron_caseta/device_trigger.py b/homeassistant/components/lutron_caseta/device_trigger.py
index 77cad154c06..b355c3dcc3f 100644
--- a/homeassistant/components/lutron_caseta/device_trigger.py
+++ b/homeassistant/components/lutron_caseta/device_trigger.py
@@ -39,6 +39,13 @@ def _reverse_dict(forward_dict: dict) -> dict:
     return {v: k for k, v in forward_dict.items()}
 
 
+LUTRON_MODEL_TO_TYPE = {
+    "RRST-W2B-XX": "SunnataKeypad_2Button",
+    "RRST-W3RL-XX": "SunnataKeypad_3ButtonRaiseLower",
+    "RRST-W4B-XX": "SunnataKeypad_4Button",
+}
+
+
 SUPPORTED_INPUTS_EVENTS_TYPES = [ACTION_PRESS, ACTION_RELEASE]
 
 LUTRON_BUTTON_TRIGGER_SCHEMA = DEVICE_TRIGGER_BASE_SCHEMA.extend(
@@ -379,9 +386,13 @@ async def async_validate_trigger_config(
     if not device:
         return config
 
-    if not (schema := DEVICE_TYPE_SCHEMA_MAP.get(device["type"])):
+    if not (
+        schema := DEVICE_TYPE_SCHEMA_MAP.get(
+            _lutron_model_to_device_type(device["model"], device["type"])
+        )
+    ):
         raise InvalidDeviceAutomationConfig(
-            f"Device type {device['type']} not supported: {config[CONF_DEVICE_ID]}"
+            f"Device model {device['model']} with type {device['type']} not supported: {config[CONF_DEVICE_ID]}"
         )
 
     return schema(config)
@@ -396,7 +407,9 @@ async def async_get_triggers(
     if not (device := get_button_device_by_dr_id(hass, device_id)):
         raise InvalidDeviceAutomationConfig(f"Device not found: {device_id}")
 
-    valid_buttons = DEVICE_TYPE_SUBTYPE_MAP_TO_LEAP.get(device["type"], {})
+    valid_buttons = DEVICE_TYPE_SUBTYPE_MAP_TO_LEAP.get(
+        _lutron_model_to_device_type(device["model"], device["type"]), {}
+    )
 
     for trigger in SUPPORTED_INPUTS_EVENTS_TYPES:
         for subtype in valid_buttons:
@@ -413,10 +426,16 @@ async def async_get_triggers(
     return triggers
 
 
-def _device_model_to_type(model: str) -> str:
+def _device_model_to_type(device_registry_model: str) -> str:
     """Convert a lutron_caseta device registry entry model to type."""
-    _, device_type = model.split(" ")
-    return device_type.replace("(", "").replace(")", "")
+    model, p_device_type = device_registry_model.split(" ")
+    device_type = p_device_type.replace("(", "").replace(")", "")
+    return _lutron_model_to_device_type(model, device_type)
+
+
+def _lutron_model_to_device_type(model: str, device_type: str) -> str:
+    """Get the mapped type based on the lutron model or type."""
+    return LUTRON_MODEL_TO_TYPE.get(model, device_type)
 
 
 async def async_attach_trigger(
diff --git a/tests/components/lutron_caseta/test_device_trigger.py b/tests/components/lutron_caseta/test_device_trigger.py
index 54cd842f0ee..b8c655a23bd 100644
--- a/tests/components/lutron_caseta/test_device_trigger.py
+++ b/tests/components/lutron_caseta/test_device_trigger.py
@@ -64,8 +64,8 @@ MOCK_BUTTON_DEVICES = [
             {"Number": 11},
         ],
         "leap_name": "Front Steps_Front Steps Sunnata Keypad",
-        "type": "SunnataKeypad_3ButtonRaiseLower",
-        "model": "PJ2-3BRL-GXX-X01",
+        "type": "SunnataKeypad",
+        "model": "RRST-W4B-XX",
         "serial": 43845547,
     },
 ]