From cc43a7a04530ec98a3ad0f3785a4b54ea02fd13f Mon Sep 17 00:00:00 2001
From: Emanuele Palazzetti <emanuele.palazzetti@gmail.com>
Date: Thu, 7 Mar 2024 15:46:56 +0100
Subject: [PATCH] fix(device): avoid state update while arming or disarming
 (#155)

---
 .../econnect_metronet/devices.py              |  9 +++++--
 tests/test_devices.py                         | 24 +++++++++++++++++++
 2 files changed, 31 insertions(+), 2 deletions(-)

diff --git a/custom_components/econnect_metronet/devices.py b/custom_components/econnect_metronet/devices.py
index c88aeb1..a5319b0 100644
--- a/custom_components/econnect_metronet/devices.py
+++ b/custom_components/econnect_metronet/devices.py
@@ -15,7 +15,9 @@
     STATE_ALARM_ARMED_HOME,
     STATE_ALARM_ARMED_NIGHT,
     STATE_ALARM_ARMED_VACATION,
+    STATE_ALARM_ARMING,
     STATE_ALARM_DISARMED,
+    STATE_ALARM_DISARMING,
     STATE_UNAVAILABLE,
 )
 from requests.exceptions import HTTPError
@@ -221,6 +223,11 @@ def get_state(self):
         Returns:
             str: One of the predefined HA alarm states.
         """
+        # If the system is arming or disarming, return the current state
+        # to prevent the state from being updated while the system is in transition.
+        if self.state in [STATE_ALARM_ARMING, STATE_ALARM_DISARMING]:
+            return self.state
+
         sectors_armed = dict(self.items(q.SECTORS, status=True))
         if not sectors_armed:
             return STATE_ALARM_DISARMED
@@ -340,7 +347,6 @@ def arm(self, code, sectors=None):
 
             with self._connection.lock(code, user_id=user_id):
                 self._connection.arm(sectors=sectors)
-                self.state = STATE_ALARM_ARMED_AWAY
         except HTTPError as err:
             _LOGGER.error(f"Device | Error while arming the system: {err.response.text}")
             raise err
@@ -368,7 +374,6 @@ def disarm(self, code, sectors=None):
 
             with self._connection.lock(code, user_id=user_id):
                 self._connection.disarm(sectors=sectors)
-                self.state = STATE_ALARM_DISARMED
         except HTTPError as err:
             _LOGGER.error(f"Device | Error while disarming the system: {err.response.text}")
             raise err
diff --git a/tests/test_devices.py b/tests/test_devices.py
index 7d9b70d..f634122 100644
--- a/tests/test_devices.py
+++ b/tests/test_devices.py
@@ -7,7 +7,9 @@
     STATE_ALARM_ARMED_HOME,
     STATE_ALARM_ARMED_NIGHT,
     STATE_ALARM_ARMED_VACATION,
+    STATE_ALARM_ARMING,
     STATE_ALARM_DISARMED,
+    STATE_ALARM_DISARMING,
     STATE_UNAVAILABLE,
 )
 from requests.exceptions import HTTPError
@@ -1604,6 +1606,28 @@ def test_get_state_armed_away_with_config(alarm_device):
     assert alarm_device.get_state() == STATE_ALARM_ARMED_AWAY
 
 
+def test_get_state_while_disarming(alarm_device):
+    # Ensure that the state is not changed while disarming
+    # Regression test for: https://github.com/palazzem/ha-econnect-alarm/issues/154
+    alarm_device._sectors_home = []
+    alarm_device._sectors_night = []
+    alarm_device._inventory = {9: {}}
+    alarm_device.state = STATE_ALARM_DISARMING
+    # Test
+    assert alarm_device.get_state() == STATE_ALARM_DISARMING
+
+
+def test_get_state_while_arming(alarm_device):
+    # Ensure that the state is not changed while arming
+    # Regression test for: https://github.com/palazzem/ha-econnect-alarm/issues/154
+    alarm_device._sectors_home = []
+    alarm_device._sectors_night = []
+    alarm_device._inventory = {9: {}}
+    alarm_device.state = STATE_ALARM_ARMING
+    # Test
+    assert alarm_device.get_state() == STATE_ALARM_ARMING
+
+
 class TestTurnOff:
     def test_required_authentication(self, alarm_device, mocker, caplog):
         # Ensure that API calls are not made when the output requires authentication