From 8a7ca822b6a7d1b7a1d2d39fd78f4308b2857da3 Mon Sep 17 00:00:00 2001
From: Nathan Spencer <natekspencer@gmail.com>
Date: Tue, 13 Dec 2022 08:41:50 -0700
Subject: [PATCH] Fix get_armed_state warning and adjust armed state map

---
 .../vivint/alarm_control_panel.py             | 49 +++++++++----------
 custom_components/vivint/manifest.json        |  2 +-
 2 files changed, 25 insertions(+), 26 deletions(-)

diff --git a/custom_components/vivint/alarm_control_panel.py b/custom_components/vivint/alarm_control_panel.py
index e7e3908..02be91a 100644
--- a/custom_components/vivint/alarm_control_panel.py
+++ b/custom_components/vivint/alarm_control_panel.py
@@ -3,16 +3,17 @@
 
 from vivintpy.enums import ArmedState
 
-from homeassistant.components.alarm_control_panel import AlarmControlPanelEntity
-from homeassistant.components.alarm_control_panel.const import (
-    SUPPORT_ALARM_ARM_AWAY,
-    SUPPORT_ALARM_ARM_HOME,
+from homeassistant.components.alarm_control_panel import (
+    AlarmControlPanelEntity,
+    AlarmControlPanelEntityFeature,
 )
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import (
     STATE_ALARM_ARMED_AWAY,
     STATE_ALARM_ARMED_HOME,
+    STATE_ALARM_ARMING,
     STATE_ALARM_DISARMED,
+    STATE_ALARM_PENDING,
     STATE_ALARM_TRIGGERED,
 )
 from homeassistant.core import HomeAssistant
@@ -22,6 +23,20 @@
 from .const import DOMAIN
 from .hub import VivintEntity
 
+ARMED_STATE_MAP = {
+    ArmedState.DISARMED: STATE_ALARM_DISARMED,
+    ArmedState.ARMING_AWAY_IN_EXIT_DELAY: STATE_ALARM_ARMING,
+    ArmedState.ARMING_STAY_IN_EXIT_DELAY: STATE_ALARM_ARMING,
+    ArmedState.ARMED_STAY: STATE_ALARM_ARMED_HOME,
+    ArmedState.ARMED_AWAY: STATE_ALARM_ARMED_AWAY,
+    ArmedState.ARMED_STAY_IN_ENTRY_DELAY: STATE_ALARM_PENDING,
+    ArmedState.ARMED_AWAY_IN_ENTRY_DELAY: STATE_ALARM_PENDING,
+    ArmedState.ALARM: STATE_ALARM_TRIGGERED,
+    ArmedState.ALARM_FIRE: STATE_ALARM_TRIGGERED,
+    ArmedState.DISABLED: STATE_ALARM_DISARMED,
+    ArmedState.WALK_TEST: STATE_ALARM_DISARMED,
+}
+
 
 async def async_setup_entry(
     hass: HomeAssistant,
@@ -47,7 +62,10 @@ class VivintAlarmControlPanelEntity(VivintEntity, AlarmControlPanelEntity):
 
     _attr_changed_by = None
     _attr_code_arm_required = False
-    _attr_supported_features = SUPPORT_ALARM_ARM_HOME | SUPPORT_ALARM_ARM_AWAY
+    _attr_supported_features = (
+        AlarmControlPanelEntityFeature.ARM_HOME
+        | AlarmControlPanelEntityFeature.ARM_AWAY
+    )
 
     @property
     def unique_id(self) -> str:
@@ -57,26 +75,7 @@ def unique_id(self) -> str:
     @property
     def state(self) -> StateType:
         """Return the state of the alarm control panel."""
-        arm_state = self.device.get_armed_state()
-        if arm_state == ArmedState.DISARMED:
-            state = STATE_ALARM_DISARMED
-        elif arm_state in [
-            ArmedState.ARMED_STAY,
-            ArmedState.ARMING_STAY_IN_EXIT_DELAY,
-            ArmedState.ARMED_STAY_IN_ENTRY_DELAY,
-        ]:
-            state = STATE_ALARM_ARMED_HOME
-        elif arm_state in [
-            ArmedState.ARMED_AWAY,
-            ArmedState.ARMING_AWAY_IN_EXIT_DELAY,
-            ArmedState.ARMED_AWAY_IN_ENTRY_DELAY,
-        ]:
-            state = STATE_ALARM_ARMED_AWAY
-        elif arm_state in [ArmedState.ALARM, ArmedState.ALARM_FIRE]:
-            state = STATE_ALARM_TRIGGERED
-        else:
-            state = None
-        return state
+        return ARMED_STATE_MAP.get(self.device.state)
 
     async def async_alarm_disarm(self, code: str | None = None) -> None:
         """Send disarm command."""
diff --git a/custom_components/vivint/manifest.json b/custom_components/vivint/manifest.json
index bf6d2af..980bfbe 100644
--- a/custom_components/vivint/manifest.json
+++ b/custom_components/vivint/manifest.json
@@ -1,7 +1,7 @@
 {
   "domain": "vivint",
   "name": "Vivint",
-  "version": "2022.12.2",
+  "version": "2022.12.3",
   "config_flow": true,
   "documentation": "https://github.com/natekspencer/hacs-vivint",
   "issue_tracker": "https://github.com/natekspencer/hacs-vivint/issues",