Improve lock handling in Yale Smart Living (#124245)

* Improve handling of locks in yalesmartalarm

* requirements

* fix coordinator setup

* Fix lock iteration

* Fix tests

* Fix review comments

homeassistant/components/yale_smart_alarm/coordinator.py

@@ -3,8 +3,9 @@
 from __future__ import annotations
 
 from datetime import timedelta
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
+from yalesmartalarmclient import YaleLock
 from yalesmartalarmclient.client import YaleSmartAlarmClient
 from yalesmartalarmclient.exceptions import AuthenticationError
 
@@ -32,6 +33,7 @@ class YaleDataUpdateCoordinator(DataUpdateCoordinator[dict[str, Any]]):
             update_interval=timedelta(seconds=DEFAULT_SCAN_INTERVAL),
             always_update=False,
         )
+        self.locks: list[YaleLock] = []
 
     async def _async_setup(self) -> None:
         """Set up connection to Yale."""
@@ -41,6 +43,7 @@ class YaleDataUpdateCoordinator(DataUpdateCoordinator[dict[str, Any]]):
                 self.entry.data[CONF_USERNAME],
                 self.entry.data[CONF_PASSWORD],
             )
+            self.locks = await self.hass.async_add_executor_job(self.yale.get_locks)
         except AuthenticationError as error:
             raise ConfigEntryAuthFailed from error
         except YALE_BASE_ERRORS as error:
@@ -51,65 +54,11 @@ class YaleDataUpdateCoordinator(DataUpdateCoordinator[dict[str, Any]]):
 
         updates = await self.hass.async_add_executor_job(self.get_updates)
 
-        locks = []
         door_windows = []
         temp_sensors = []
 
         for device in updates["cycle"]["device_status"]:
             state = device["status1"]
-            if device["type"] == "device_type.door_lock":
-                lock_status_str = device["minigw_lock_status"]
-                lock_status = int(str(lock_status_str or 0), 16)
-                closed = (lock_status & 16) == 16
-                locked = (lock_status & 1) == 1
-                if not lock_status and "device_status.lock" in state:
-                    device["_state"] = "locked"
-                    device["_state2"] = "unknown"
-                    locks.append(device)
-                    continue
-                if not lock_status and "device_status.unlock" in state:
-                    device["_state"] = "unlocked"
-                    device["_state2"] = "unknown"
-                    locks.append(device)
-                    continue
-                if (
-                    lock_status
-                    and (
-                        "device_status.lock" in state or "device_status.unlock" in state
-                    )
-                    and closed
-                    and locked
-                ):
-                    device["_state"] = "locked"
-                    device["_state2"] = "closed"
-                    locks.append(device)
-                    continue
-                if (
-                    lock_status
-                    and (
-                        "device_status.lock" in state or "device_status.unlock" in state
-                    )
-                    and closed
-                    and not locked
-                ):
-                    device["_state"] = "unlocked"
-                    device["_state2"] = "closed"
-                    locks.append(device)
-                    continue
-                if (
-                    lock_status
-                    and (
-                        "device_status.lock" in state or "device_status.unlock" in state
-                    )
-                    and not closed
-                ):
-                    device["_state"] = "unlocked"
-                    device["_state2"] = "open"
-                    locks.append(device)
-                    continue
-                device["_state"] = "unavailable"
-                locks.append(device)
-                continue
             if device["type"] == "device_type.door_contact":
                 if "device_status.dc_close" in state:
                     device["_state"] = "closed"
@@ -128,19 +77,16 @@ class YaleDataUpdateCoordinator(DataUpdateCoordinator[dict[str, Any]]):
         _sensor_map = {
             contact["address"]: contact["_state"] for contact in door_windows
         }
-        _lock_map = {lock["address"]: lock["_state"] for lock in locks}
         _temp_map = {temp["address"]: temp["status_temp"] for temp in temp_sensors}
 
         return {
             "alarm": updates["arm_status"],
-            "locks": locks,
             "door_windows": door_windows,
             "temp_sensors": temp_sensors,
             "status": updates["status"],
             "online": updates["online"],
             "sensor_map": _sensor_map,
             "temp_map": _temp_map,
-            "lock_map": _lock_map,
             "panel_info": updates["panel_info"],
         }
 
@@ -149,6 +95,13 @@ class YaleDataUpdateCoordinator(DataUpdateCoordinator[dict[str, Any]]):
         try:
             arm_status = self.yale.get_armed_status()
             data = self.yale.get_information()
+            if TYPE_CHECKING:
+                assert data.cycle
+            for device in data.cycle["data"]["device_status"]:
+                if device["type"] == YaleLock.DEVICE_TYPE:
+                    for lock in self.locks:
+                        if lock.name == device["name"]:
+                            lock.update(device)
         except AuthenticationError as error:
             raise ConfigEntryAuthFailed from error
         except YALE_BASE_ERRORS as error:

homeassistant/components/yale_smart_alarm/entity.py

@@ -1,5 +1,7 @@
 """Base class for yale_smart_alarm entity."""
 
+from yalesmartalarmclient import YaleLock
+
 from homeassistant.const import CONF_NAME, CONF_USERNAME
 from homeassistant.helpers.device_registry import CONNECTION_NETWORK_MAC, DeviceInfo
 from homeassistant.helpers.entity import Entity
@@ -9,7 +11,7 @@ from .const import DOMAIN, MANUFACTURER, MODEL
 from .coordinator import YaleDataUpdateCoordinator
 
 
-class YaleEntity(CoordinatorEntity[YaleDataUpdateCoordinator], Entity):
+class YaleEntity(CoordinatorEntity[YaleDataUpdateCoordinator]):
     """Base implementation for Yale device."""
 
     _attr_has_entity_name = True
@@ -23,7 +25,25 @@ class YaleEntity(CoordinatorEntity[YaleDataUpdateCoordinator], Entity):
             manufacturer=MANUFACTURER,
             model=MODEL,
             identifiers={(DOMAIN, data["address"])},
-            via_device=(DOMAIN, self.coordinator.entry.data[CONF_USERNAME]),
+            via_device=(DOMAIN, coordinator.entry.data[CONF_USERNAME]),
+        )
+
+
+class YaleLockEntity(CoordinatorEntity[YaleDataUpdateCoordinator]):
+    """Base implementation for Yale lock device."""
+
+    _attr_has_entity_name = True
+
+    def __init__(self, coordinator: YaleDataUpdateCoordinator, lock: YaleLock) -> None:
+        """Initialize an Yale device."""
+        super().__init__(coordinator)
+        self._attr_unique_id: str = lock.sid()
+        self._attr_device_info = DeviceInfo(
+            name=lock.name,
+            manufacturer=MANUFACTURER,
+            model=MODEL,
+            identifiers={(DOMAIN, lock.sid())},
+            via_device=(DOMAIN, coordinator.entry.data[CONF_USERNAME]),
         )
 
 

homeassistant/components/yale_smart_alarm/lock.py

@@ -2,9 +2,16 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any
+from typing import Any
 
-from homeassistant.components.lock import LockEntity
+from yalesmartalarmclient import YaleLock, YaleLockState
+
+from homeassistant.components.lock import (
+    STATE_LOCKED,
+    STATE_OPEN,
+    STATE_UNLOCKED,
+    LockEntity,
+)
 from homeassistant.const import ATTR_CODE
 from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import HomeAssistantError, ServiceValidationError
@@ -18,7 +25,13 @@ from .const import (
     YALE_ALL_ERRORS,
 )
 from .coordinator import YaleDataUpdateCoordinator
-from .entity import YaleEntity
+from .entity import YaleLockEntity
+
+LOCK_STATE_MAP = {
+    YaleLockState.LOCKED: STATE_LOCKED,
+    YaleLockState.UNLOCKED: STATE_UNLOCKED,
+    YaleLockState.DOOR_OPEN: STATE_OPEN,
+}
 
 
 async def async_setup_entry(
@@ -30,68 +43,62 @@ async def async_setup_entry(
     code_format = entry.options.get(CONF_LOCK_CODE_DIGITS, DEFAULT_LOCK_CODE_DIGITS)
 
     async_add_entities(
-        YaleDoorlock(coordinator, data, code_format)
+        YaleDoorlock(coordinator, lock, code_format) for lock in coordinator.locks
-        for data in coordinator.data["locks"]
     )
 
 
-class YaleDoorlock(YaleEntity, LockEntity):
+class YaleDoorlock(YaleLockEntity, LockEntity):
     """Representation of a Yale doorlock."""
 
     _attr_name = None
 
     def __init__(
-        self, coordinator: YaleDataUpdateCoordinator, data: dict, code_format: int
+        self, coordinator: YaleDataUpdateCoordinator, lock: YaleLock, code_format: int
     ) -> None:
         """Initialize the Yale Lock Device."""
-        super().__init__(coordinator, data)
+        super().__init__(coordinator, lock)
         self._attr_code_format = rf"^\d{{{code_format}}}$"
-        self.lock_name: str = data["name"]
+        self.lock_data = lock
 
     async def async_unlock(self, **kwargs: Any) -> None:
         """Send unlock command."""
         code: str | None = kwargs.get(ATTR_CODE)
-        return await self.async_set_lock("unlocked", code)
+        return await self.async_set_lock(YaleLockState.UNLOCKED, code)
 
     async def async_lock(self, **kwargs: Any) -> None:
         """Send lock command."""
-        return await self.async_set_lock("locked", None)
+        return await self.async_set_lock(YaleLockState.LOCKED, None)
 
-    async def async_set_lock(self, command: str, code: str | None) -> None:
+    async def async_set_lock(self, state: YaleLockState, code: str | None) -> None:
         """Set lock."""
-        if TYPE_CHECKING:
+        if state is YaleLockState.UNLOCKED and not code:
-            assert self.coordinator.yale, "Connection to API is missing"
-        if command == "unlocked" and not code:
             raise ServiceValidationError(
                 translation_domain=DOMAIN,
                 translation_key="no_code",
             )
 
+        lock_state = False
         try:
-            get_lock = await self.hass.async_add_executor_job(
+            if state is YaleLockState.LOCKED:
-                self.coordinator.yale.lock_api.get, self.lock_name
-            )
-            if get_lock and command == "locked":
                 lock_state = await self.hass.async_add_executor_job(
-                    self.coordinator.yale.lock_api.close_lock,
+                    self.lock_data.close
-                    get_lock,
                 )
-            if code and get_lock and command == "unlocked":
+            if code and state is YaleLockState.UNLOCKED:
                 lock_state = await self.hass.async_add_executor_job(
-                    self.coordinator.yale.lock_api.open_lock, get_lock, code
+                    self.lock_data.open, code
                 )
         except YALE_ALL_ERRORS as error:
             raise HomeAssistantError(
                 translation_domain=DOMAIN,
                 translation_key="set_lock",
                 translation_placeholders={
-                    "name": self.lock_name,
+                    "name": self.lock_data.name,
                     "error": str(error),
                 },
             ) from error
 
         if lock_state:
-            self.coordinator.data["lock_map"][self._attr_unique_id] = command
+            self.lock_data.set_state(state)
             self.async_write_ha_state()
             return
         raise HomeAssistantError(
@@ -102,4 +109,9 @@ class YaleDoorlock(YaleEntity, LockEntity):
     @property
     def is_locked(self) -> bool | None:
         """Return true if the lock is locked."""
-        return bool(self.coordinator.data["lock_map"][self._attr_unique_id] == "locked")
+        return LOCK_STATE_MAP.get(self.lock_data.state()) == STATE_LOCKED
+
+    @property
+    def is_open(self) -> bool | None:
+        """Return true if the lock is open."""
+        return LOCK_STATE_MAP.get(self.lock_data.state()) == STATE_OPEN

homeassistant/components/yale_smart_alarm/manifest.json

@@ -6,5 +6,5 @@
   "documentation": "https://www.home-assistant.io/integrations/yale_smart_alarm",
   "iot_class": "cloud_polling",
   "loggers": ["yalesmartalarmclient"],
-  "requirements": ["yalesmartalarmclient==0.4.0"]
+  "requirements": ["yalesmartalarmclient==0.4.2"]
 }

requirements_all.txt

@@ -3002,7 +3002,7 @@ xmltodict==0.13.0
 xs1-api-client==3.0.0
 
 # homeassistant.components.yale_smart_alarm
-yalesmartalarmclient==0.4.0
+yalesmartalarmclient==0.4.2
 
 # homeassistant.components.august
 # homeassistant.components.yale

requirements_test_all.txt

@@ -2388,7 +2388,7 @@ xknxproject==3.7.1
 xmltodict==0.13.0
 
 # homeassistant.components.yale_smart_alarm
-yalesmartalarmclient==0.4.0
+yalesmartalarmclient==0.4.2
 
 # homeassistant.components.august
 # homeassistant.components.yale

tests/components/yale_smart_alarm/conftest.py

@@ -7,7 +7,7 @@ from typing import Any
 from unittest.mock import Mock, patch
 
 import pytest
-from yalesmartalarmclient import YaleSmartAlarmData
+from yalesmartalarmclient import YaleDoorManAPI, YaleLock, YaleSmartAlarmData
 from yalesmartalarmclient.const import YALE_STATE_ARM_FULL
 
 from homeassistant.components.yale_smart_alarm.const import DOMAIN, PLATFORMS
@@ -53,16 +53,28 @@ async def load_config_entry(
 
         config_entry.add_to_hass(hass)
 
+        cycle = get_data.cycle["data"]
+        data = {"data": cycle["device_status"]}
+
         with patch(
             "homeassistant.components.yale_smart_alarm.coordinator.YaleSmartAlarmClient",
             autospec=True,
         ) as mock_client_class:
             client = mock_client_class.return_value
             client.auth = Mock()
-            client.lock_api = Mock()
+            client.auth.get_authenticated = Mock(return_value=data)
+            client.auth.post_authenticated = Mock(return_value={"code": "000"})
+            client.lock_api = YaleDoorManAPI(client.auth)
+            locks = [
+                YaleLock(device, lock_api=client.lock_api)
+                for device in cycle["device_status"]
+                if device["type"] == YaleLock.DEVICE_TYPE
+            ]
+            client.get_locks.return_value = locks
             client.get_all.return_value = get_all_data
             client.get_information.return_value = get_data
             client.get_armed_status.return_value = YALE_STATE_ARM_FULL
+
             await hass.config_entries.async_setup(config_entry.entry_id)
             await hass.async_block_till_done()
 
@@ -78,7 +90,7 @@ def get_fixture_data() -> dict[str, Any]:
     return json_data
 
 
-@pytest.fixture(name="get_data", scope="package")
+@pytest.fixture(name="get_data")
 def get_update_data(loaded_fixture: dict[str, Any]) -> YaleSmartAlarmData:
     """Load update data and return."""
 
@@ -94,7 +106,7 @@ def get_update_data(loaded_fixture: dict[str, Any]) -> YaleSmartAlarmData:
     )
 
 
-@pytest.fixture(name="get_all_data", scope="package")
+@pytest.fixture(name="get_all_data")
 def get_diag_data(loaded_fixture: dict[str, Any]) -> YaleSmartAlarmData:
     """Load all data and return."""
 

tests/components/yale_smart_alarm/snapshots/test_diagnostics.ambr

@@ -24,8 +24,6 @@
         'capture_latest': None,
         'device_status': list([
           dict({
-            '_state': 'locked',
-            '_state2': 'closed',
             'address': '**REDACTED**',
             'area': '1',
             'bypass': '0',
@@ -86,8 +84,6 @@
             'type_no': '72',
           }),
           dict({
-            '_state': 'unlocked',
-            '_state2': 'unknown',
             'address': '**REDACTED**',
             'area': '1',
             'bypass': '0',
@@ -147,8 +143,6 @@
             'type_no': '72',
           }),
           dict({
-            '_state': 'locked',
-            '_state2': 'unknown',
             'address': '**REDACTED**',
             'area': '1',
             'bypass': '0',
@@ -391,8 +385,6 @@
             'type_no': '4',
           }),
           dict({
-            '_state': 'unlocked',
-            '_state2': 'closed',
             'address': '**REDACTED**',
             'area': '1',
             'bypass': '0',
@@ -453,8 +445,6 @@
             'type_no': '72',
           }),
           dict({
-            '_state': 'unlocked',
-            '_state2': 'open',
             'address': '**REDACTED**',
             'area': '1',
             'bypass': '0',
@@ -515,7 +505,6 @@
             'type_no': '72',
           }),
           dict({
-            '_state': 'unavailable',
             'address': '**REDACTED**',
             'area': '1',
             'bypass': '0',

tests/components/yale_smart_alarm/snapshots/test_lock.ambr

@@ -236,7 +236,7 @@
     'last_changed': <ANY>,
     'last_reported': <ANY>,
     'last_updated': <ANY>,
-    'state': 'unlocked',
+    'state': 'open',
   })
 # ---
 # name: test_lock[load_platforms0][lock.device9-entry]

tests/components/yale_smart_alarm/test_lock.py

@@ -47,8 +47,6 @@ async def test_lock_service_calls(
     hass: HomeAssistant,
     get_data: YaleSmartAlarmData,
     load_config_entry: tuple[MockConfigEntry, Mock],
-    entity_registry: er.EntityRegistry,
-    snapshot: SnapshotAssertion,
 ) -> None:
     """Test the Yale Smart Alarm lock."""
 
@@ -101,8 +99,6 @@ async def test_lock_service_call_fails(
     hass: HomeAssistant,
     get_data: YaleSmartAlarmData,
     load_config_entry: tuple[MockConfigEntry, Mock],
-    entity_registry: er.EntityRegistry,
-    snapshot: SnapshotAssertion,
 ) -> None:
     """Test the Yale Smart Alarm lock service call fails."""
 
@@ -153,8 +149,6 @@ async def test_lock_service_call_fails_with_incorrect_status(
     hass: HomeAssistant,
     get_data: YaleSmartAlarmData,
     load_config_entry: tuple[MockConfigEntry, Mock],
-    entity_registry: er.EntityRegistry,
-    snapshot: SnapshotAssertion,
 ) -> None:
     """Test the Yale Smart Alarm lock service call fails with incorrect return state."""
 
@@ -163,9 +157,7 @@ async def test_lock_service_call_fails_with_incorrect_status(
     data = deepcopy(get_data.cycle)
     data["data"] = data["data"].pop("device_status")
 
-    client.auth.get_authenticated = Mock(return_value=data)
     client.auth.post_authenticated = Mock(return_value={"code": "FFF"})
-    client.lock_api = YaleDoorManAPI(client.auth)
 
     state = hass.states.get("lock.device1")
     assert state.state == "locked"