From 0ca4314d486e14d23b50cd8c36b615596e9ba003 Mon Sep 17 00:00:00 2001
From: tronikos <tronikos@users.noreply.github.com>
Date: Sat, 8 Jun 2024 13:54:32 -0700
Subject: [PATCH] Make supported_features of manual alarm_control_panel
 configurable (#119122)

---
 .../components/manual/alarm_control_panel.py  | 26 ++++---
 .../manual/test_alarm_control_panel.py        | 68 +++++++++++++++++++
 2 files changed, 86 insertions(+), 8 deletions(-)

diff --git a/homeassistant/components/manual/alarm_control_panel.py b/homeassistant/components/manual/alarm_control_panel.py
index 37580011a5e..5b344dd01ac 100644
--- a/homeassistant/components/manual/alarm_control_panel.py
+++ b/homeassistant/components/manual/alarm_control_panel.py
@@ -42,6 +42,7 @@ import homeassistant.util.dt as dt_util
 
 _LOGGER = logging.getLogger(__name__)
 
+CONF_ARMING_STATES = "arming_states"
 CONF_CODE_TEMPLATE = "code_template"
 CONF_CODE_ARM_REQUIRED = "code_arm_required"
 
@@ -71,6 +72,14 @@ SUPPORTED_ARMING_STATES = [
     if state not in (STATE_ALARM_DISARMED, STATE_ALARM_TRIGGERED)
 ]
 
+SUPPORTED_ARMING_STATE_TO_FEATURE = {
+    STATE_ALARM_ARMED_AWAY: AlarmControlPanelEntityFeature.ARM_AWAY,
+    STATE_ALARM_ARMED_HOME: AlarmControlPanelEntityFeature.ARM_HOME,
+    STATE_ALARM_ARMED_NIGHT: AlarmControlPanelEntityFeature.ARM_NIGHT,
+    STATE_ALARM_ARMED_VACATION: AlarmControlPanelEntityFeature.ARM_VACATION,
+    STATE_ALARM_ARMED_CUSTOM_BYPASS: AlarmControlPanelEntityFeature.ARM_CUSTOM_BYPASS,
+}
+
 ATTR_PREVIOUS_STATE = "previous_state"
 ATTR_NEXT_STATE = "next_state"
 
@@ -128,6 +137,9 @@ PLATFORM_SCHEMA = vol.Schema(
             vol.Optional(
                 CONF_DISARM_AFTER_TRIGGER, default=DEFAULT_DISARM_AFTER_TRIGGER
             ): cv.boolean,
+            vol.Optional(CONF_ARMING_STATES, default=SUPPORTED_ARMING_STATES): vol.All(
+                cv.ensure_list, [vol.In(SUPPORTED_ARMING_STATES)]
+            ),
             vol.Optional(STATE_ALARM_ARMED_AWAY, default={}): _state_schema(
                 STATE_ALARM_ARMED_AWAY
             ),
@@ -188,14 +200,6 @@ class ManualAlarm(AlarmControlPanelEntity, RestoreEntity):
     """
 
     _attr_should_poll = False
-    _attr_supported_features = (
-        AlarmControlPanelEntityFeature.ARM_HOME
-        | AlarmControlPanelEntityFeature.ARM_AWAY
-        | AlarmControlPanelEntityFeature.ARM_NIGHT
-        | AlarmControlPanelEntityFeature.ARM_VACATION
-        | AlarmControlPanelEntityFeature.TRIGGER
-        | AlarmControlPanelEntityFeature.ARM_CUSTOM_BYPASS
-    )
 
     def __init__(
         self,
@@ -233,6 +237,12 @@ class ManualAlarm(AlarmControlPanelEntity, RestoreEntity):
             state: config[state][CONF_ARMING_TIME] for state in SUPPORTED_ARMING_STATES
         }
 
+        self._attr_supported_features = AlarmControlPanelEntityFeature.TRIGGER
+        for arming_state in config.get(CONF_ARMING_STATES, SUPPORTED_ARMING_STATES):
+            self._attr_supported_features |= SUPPORTED_ARMING_STATE_TO_FEATURE[
+                arming_state
+            ]
+
     @property
     def state(self) -> str:
         """Return the state of the device."""
diff --git a/tests/components/manual/test_alarm_control_panel.py b/tests/components/manual/test_alarm_control_panel.py
index 5910cc3ec9b..6c9ba9ee9a0 100644
--- a/tests/components/manual/test_alarm_control_panel.py
+++ b/tests/components/manual/test_alarm_control_panel.py
@@ -7,6 +7,7 @@ from freezegun import freeze_time
 import pytest
 
 from homeassistant.components import alarm_control_panel
+from homeassistant.components.alarm_control_panel import AlarmControlPanelEntityFeature
 from homeassistant.components.demo import alarm_control_panel as demo
 from homeassistant.const import (
     ATTR_CODE,
@@ -1456,3 +1457,70 @@ async def test_restore_state_triggered_long_ago(hass: HomeAssistant) -> None:
 
     state = hass.states.get(entity_id)
     assert state.state == STATE_ALARM_DISARMED
+
+
+async def test_default_arming_states(hass: HomeAssistant) -> None:
+    """Test default arming_states."""
+    assert await async_setup_component(
+        hass,
+        alarm_control_panel.DOMAIN,
+        {
+            "alarm_control_panel": {
+                "platform": "manual",
+                "name": "test",
+            }
+        },
+    )
+    await hass.async_block_till_done()
+
+    state = hass.states.get("alarm_control_panel.test")
+    assert state.attributes["supported_features"] == (
+        AlarmControlPanelEntityFeature.ARM_HOME
+        | AlarmControlPanelEntityFeature.ARM_AWAY
+        | AlarmControlPanelEntityFeature.ARM_NIGHT
+        | AlarmControlPanelEntityFeature.ARM_VACATION
+        | AlarmControlPanelEntityFeature.TRIGGER
+        | AlarmControlPanelEntityFeature.ARM_CUSTOM_BYPASS
+    )
+
+
+async def test_arming_states(hass: HomeAssistant) -> None:
+    """Test arming_states."""
+    assert await async_setup_component(
+        hass,
+        alarm_control_panel.DOMAIN,
+        {
+            "alarm_control_panel": {
+                "platform": "manual",
+                "name": "test",
+                "arming_states": ["armed_away", "armed_home"],
+            }
+        },
+    )
+    await hass.async_block_till_done()
+
+    state = hass.states.get("alarm_control_panel.test")
+    assert state.attributes["supported_features"] == (
+        AlarmControlPanelEntityFeature.ARM_HOME
+        | AlarmControlPanelEntityFeature.ARM_AWAY
+        | AlarmControlPanelEntityFeature.TRIGGER
+    )
+
+
+async def test_invalid_arming_states(hass: HomeAssistant) -> None:
+    """Test invalid arming_states."""
+    assert await async_setup_component(
+        hass,
+        alarm_control_panel.DOMAIN,
+        {
+            "alarm_control_panel": {
+                "platform": "manual",
+                "name": "test",
+                "arming_states": ["invalid", "armed_home"],
+            }
+        },
+    )
+    await hass.async_block_till_done()
+
+    state = hass.states.get("alarm_control_panel.test")
+    assert state is None