From f9e9202e2de9d2b418b0ee3b5860ef9106136685 Mon Sep 17 00:00:00 2001
From: Erik Montnemery <erik@montnemery.com>
Date: Sun, 13 Jun 2021 00:07:25 +0200
Subject: [PATCH] Improve editing of device triggers referencing non-added
 alarm (#51701)

---
 .../alarm_control_panel/device_trigger.py     |  10 +-
 .../test_device_trigger.py                    | 102 ++++++++++--------
 2 files changed, 62 insertions(+), 50 deletions(-)

diff --git a/homeassistant/components/alarm_control_panel/device_trigger.py b/homeassistant/components/alarm_control_panel/device_trigger.py
index cae9161abf9..f89e03e7326 100644
--- a/homeassistant/components/alarm_control_panel/device_trigger.py
+++ b/homeassistant/components/alarm_control_panel/device_trigger.py
@@ -14,7 +14,6 @@ from homeassistant.components.automation import AutomationActionType
 from homeassistant.components.device_automation import DEVICE_TRIGGER_BASE_SCHEMA
 from homeassistant.components.homeassistant.triggers import state as state_trigger
 from homeassistant.const import (
-    ATTR_SUPPORTED_FEATURES,
     CONF_DEVICE_ID,
     CONF_DOMAIN,
     CONF_ENTITY_ID,
@@ -30,6 +29,7 @@ from homeassistant.const import (
 )
 from homeassistant.core import CALLBACK_TYPE, HomeAssistant
 from homeassistant.helpers import config_validation as cv, entity_registry
+from homeassistant.helpers.entity import get_supported_features
 from homeassistant.helpers.typing import ConfigType
 
 from . import DOMAIN
@@ -62,13 +62,7 @@ async def async_get_triggers(
         if entry.domain != DOMAIN:
             continue
 
-        entity_state = hass.states.get(entry.entity_id)
-
-        # We need a state or else we can't populate the HVAC and preset modes.
-        if entity_state is None:
-            continue
-
-        supported_features = entity_state.attributes[ATTR_SUPPORTED_FEATURES]
+        supported_features = get_supported_features(hass, entry.entity_id)
 
         # Add triggers for each entity that belongs to this integration
         base_trigger = {
diff --git a/tests/components/alarm_control_panel/test_device_trigger.py b/tests/components/alarm_control_panel/test_device_trigger.py
index 3380cdd9654..8859915b911 100644
--- a/tests/components/alarm_control_panel/test_device_trigger.py
+++ b/tests/components/alarm_control_panel/test_device_trigger.py
@@ -48,7 +48,48 @@ def calls(hass):
     return async_mock_service(hass, "test", "automation")
 
 
-async def test_get_triggers(hass, device_reg, entity_reg):
+@pytest.mark.parametrize(
+    "set_state,features_reg,features_state,expected_trigger_types",
+    [
+        (False, 0, 0, ["triggered", "disarmed", "arming"]),
+        (
+            False,
+            15,
+            0,
+            [
+                "triggered",
+                "disarmed",
+                "arming",
+                "armed_home",
+                "armed_away",
+                "armed_night",
+            ],
+        ),
+        (True, 0, 0, ["triggered", "disarmed", "arming"]),
+        (
+            True,
+            0,
+            15,
+            [
+                "triggered",
+                "disarmed",
+                "arming",
+                "armed_home",
+                "armed_away",
+                "armed_night",
+            ],
+        ),
+    ],
+)
+async def test_get_triggers(
+    hass,
+    device_reg,
+    entity_reg,
+    set_state,
+    features_reg,
+    features_state,
+    expected_trigger_types,
+):
     """Test we get the expected triggers from an alarm_control_panel."""
     config_entry = MockConfigEntry(domain="test", data={})
     config_entry.add_to_hass(hass)
@@ -56,53 +97,30 @@ async def test_get_triggers(hass, device_reg, entity_reg):
         config_entry_id=config_entry.entry_id,
         connections={(device_registry.CONNECTION_NETWORK_MAC, "12:34:56:AB:CD:EF")},
     )
-    entity_reg.async_get_or_create(DOMAIN, "test", "5678", device_id=device_entry.id)
-    hass.states.async_set(
-        "alarm_control_panel.test_5678", "attributes", {"supported_features": 15}
+    entity_reg.async_get_or_create(
+        DOMAIN,
+        "test",
+        "5678",
+        device_id=device_entry.id,
+        supported_features=features_reg,
     )
-    expected_triggers = [
+    if set_state:
+        hass.states.async_set(
+            "alarm_control_panel.test_5678",
+            "attributes",
+            {"supported_features": features_state},
+        )
+    expected_triggers = []
+
+    expected_triggers += [
         {
             "platform": "device",
             "domain": DOMAIN,
-            "type": "disarmed",
+            "type": trigger,
             "device_id": device_entry.id,
             "entity_id": f"{DOMAIN}.test_5678",
-        },
-        {
-            "platform": "device",
-            "domain": DOMAIN,
-            "type": "triggered",
-            "device_id": device_entry.id,
-            "entity_id": f"{DOMAIN}.test_5678",
-        },
-        {
-            "platform": "device",
-            "domain": DOMAIN,
-            "type": "arming",
-            "device_id": device_entry.id,
-            "entity_id": f"{DOMAIN}.test_5678",
-        },
-        {
-            "platform": "device",
-            "domain": DOMAIN,
-            "type": "armed_home",
-            "device_id": device_entry.id,
-            "entity_id": f"{DOMAIN}.test_5678",
-        },
-        {
-            "platform": "device",
-            "domain": DOMAIN,
-            "type": "armed_away",
-            "device_id": device_entry.id,
-            "entity_id": f"{DOMAIN}.test_5678",
-        },
-        {
-            "platform": "device",
-            "domain": DOMAIN,
-            "type": "armed_night",
-            "device_id": device_entry.id,
-            "entity_id": f"{DOMAIN}.test_5678",
-        },
+        }
+        for trigger in expected_trigger_types
     ]
     triggers = await async_get_device_automations(hass, "trigger", device_entry.id)
     assert_lists_same(triggers, expected_triggers)