From 61a9c994dce6b547ee1af4bcc9395c07c66d9453 Mon Sep 17 00:00:00 2001
From: puddly <32534428+puddly@users.noreply.github.com>
Date: Sat, 23 Dec 2023 00:36:11 -0500
Subject: [PATCH] Implement command priority

---
 pyproject.toml      |  2 +-
 zigpy_zigate/api.py | 13 +++++++++++--
 2 files changed, 12 insertions(+), 3 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index 7d76f7e..a9a15b8 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -15,7 +15,7 @@ license = {text = "GPL-3.0"}
 requires-python = ">=3.8"
 dependencies = [
     "voluptuous",
-    "zigpy>=0.60.0",
+    "zigpy>=0.60.2",
     "pyusb>=1.1.0",
     "gpiozero",
     'async-timeout; python_version<"3.11"',
diff --git a/zigpy_zigate/api.py b/zigpy_zigate/api.py
index 4556579..34ee054 100644
--- a/zigpy_zigate/api.py
+++ b/zigpy_zigate/api.py
@@ -6,6 +6,7 @@
 import logging
 from typing import Any, Dict
 
+from zigpy.datastructures import PriorityLock
 import zigpy.exceptions
 
 import zigpy_zigate.uart
@@ -217,7 +218,7 @@ def __init__(self, device_config: Dict[str, Any]):
         self._uart = None
         self._awaiting = {}
         self._status_awaiting = {}
-        self._lock = asyncio.Lock()
+        self._lock = PriorityLock()
 
         self.network_state = None
 
@@ -295,6 +296,14 @@ async def wait_for_response(self, wait_response):
                 self._awaiting[wait_response].cancel()
                 del self._awaiting[wait_response]
 
+    def _get_command_priority(self, cmd):
+        return {
+            # Watchdog command is prioritized
+            CommandId.SET_TIMESERVER: 9999,
+            # APS command is deprioritized
+            CommandId.SEND_RAW_APS_DATA_PACKET: -1,
+        }.get(cmd, 0)
+
     async def command(
         self,
         cmd,
@@ -303,7 +312,7 @@ async def command(
         wait_status=True,
         timeout=COMMAND_TIMEOUT,
     ):
-        async with self._lock:
+        async with self._lock(priority=self._get_command_priority(cmd)):
             tries = 3
 
             tasks = []