Merge branch 'master' into version_12

home-assistant-libs · Sep 20, 2023 · 8ff929b · 8ff929b
2 parents 6209276 + 7c6e5a6
commit 8ff929b
Show file tree

Hide file tree

Showing 6 changed files with 46 additions and 24 deletions.
diff --git a/test/conftest.py b/test/conftest.py
@@ -161,9 +161,11 @@ async def receive():
         """Return a websocket message."""
         await asyncio.sleep(0)
 
-        message = messages.popleft()
-        if not messages:
+        try:
+            message = messages.popleft()
+        except IndexError:
             ws_client.closed = True
+            return WSMessage(WSMsgType.CLOSED, None, None)
 
         return message
 

diff --git a/test/test_client.py b/test/test_client.py
@@ -461,3 +461,18 @@ async def test_additional_user_agent_components(client_session, url):
                 },
             }
         )
+
+
+async def test_pop_future_none(client_session, url, driver_ready):
+    """Test when a future has been cleared from futures dict, popping still works."""
+    client = Client(url, client_session)
+    await client.connect()
+
+    assert client.connected
+
+    asyncio.create_task(client.listen(driver_ready))
+
+    await driver_ready.wait()
+
+    with pytest.raises(asyncio.CancelledError):
+        await client.async_send_command({"command": "some_command"})
diff --git a/test/test_main.py b/test/test_main.py
@@ -72,4 +72,4 @@ def test_connect(client_session, url, ws_client):
         main()
 
     assert sys_exit.value.code == 0
-    assert ws_client.receive.call_count == 4
+    assert ws_client.receive.call_count == 5
diff --git a/zwave_js_server/client.py b/zwave_js_server/client.py
@@ -115,7 +115,7 @@ async def async_send_command(
         try:
             return await future
         finally:
-            self._result_futures.pop(message_id)
+            self._result_futures.pop(message_id, None)
 
     async def async_send_command_no_wait(
         self, message: dict[str, Any], require_schema: int | None = None

diff --git a/zwave_js_server/model/node/__init__.py b/zwave_js_server/model/node/__init__.py
@@ -1058,25 +1058,27 @@ def handle_metadata_updated(self, event: Event) -> None:
 
     def handle_notification(self, event: Event) -> None:
         """Process a node notification event."""
-        command_class = CommandClass(event.data["ccId"])
-        if command_class == CommandClass.NOTIFICATION:
-            event.data["notification"] = NotificationNotification(
-                self, cast(NotificationNotificationDataType, event.data)
-            )
-        elif command_class == CommandClass.SWITCH_MULTILEVEL:
-            event.data["notification"] = MultilevelSwitchNotification(
-                self, cast(MultilevelSwitchNotificationDataType, event.data)
-            )
-        elif command_class == CommandClass.ENTRY_CONTROL:
-            event.data["notification"] = EntryControlNotification(
-                self, cast(EntryControlNotificationDataType, event.data)
-            )
-        elif command_class == CommandClass.POWERLEVEL:
-            event.data["notification"] = PowerLevelNotification(
-                self, cast(PowerLevelNotificationDataType, event.data)
-            )
-        else:
-            _LOGGER.info("Unhandled notification command class: %s", command_class.name)
+        match command_class := CommandClass(event.data["ccId"]):
+            case CommandClass.NOTIFICATION:
+                event.data["notification"] = NotificationNotification(
+                    self, cast(NotificationNotificationDataType, event.data)
+                )
+            case CommandClass.SWITCH_MULTILEVEL:
+                event.data["notification"] = MultilevelSwitchNotification(
+                    self, cast(MultilevelSwitchNotificationDataType, event.data)
+                )
+            case CommandClass.ENTRY_CONTROL:
+                event.data["notification"] = EntryControlNotification(
+                    self, cast(EntryControlNotificationDataType, event.data)
+                )
+            case CommandClass.POWERLEVEL:
+                event.data["notification"] = PowerLevelNotification(
+                    self, cast(PowerLevelNotificationDataType, event.data)
+                )
+            case _:
+                _LOGGER.info(
+                    "Unhandled notification command class: %s", command_class.name
+                )
 
     def handle_firmware_update_progress(self, event: Event) -> None:
         """Process a node firmware update progress event."""

diff --git a/zwave_js_server/model/node/event_model.py b/zwave_js_server/model/node/event_model.py
@@ -7,6 +7,7 @@
 from ...event import BaseEventModel
 from ..notification import (
     EntryControlNotificationArgsDataType,
+    MultilevelSwitchNotificationArgsDataType,
     NotificationNotificationArgsDataType,
     PowerLevelNotificationArgsDataType,
 )
@@ -106,12 +107,14 @@ class NotificationEventModel(BaseNodeEventModel):
     """Model for `notification` event data."""
 
     event: Literal["notification"]
-    endpoint: int
+    endpointIndex: int
+    nodeId: int
     ccId: CommandClass
     args: (
         NotificationNotificationArgsDataType
         | EntryControlNotificationArgsDataType
         | PowerLevelNotificationArgsDataType
+        | MultilevelSwitchNotificationArgsDataType
     )