From a14912dec1af0105afe46a3a480feea21ec5030d Mon Sep 17 00:00:00 2001
From: puddly <32534428+puddly@users.noreply.github.com>
Date: Fri, 10 Nov 2023 11:37:44 -0500
Subject: [PATCH] Unit test request locking

---
 tests/test_api.py   | 114 ++++++++++++++++++++++++++++++++------------
 zigpy_deconz/api.py |   7 ++-
 2 files changed, 87 insertions(+), 34 deletions(-)

diff --git a/tests/test_api.py b/tests/test_api.py
index e8773d1..93997bb 100644
--- a/tests/test_api.py
+++ b/tests/test_api.py
@@ -135,6 +135,35 @@ def receiver(data):
     return inner
 
 
+def send_network_state(
+    api,
+    network_state: deconz_api.NetworkState2 = deconz_api.NetworkState2.CONNECTED,
+    device_state: deconz_api.DeviceStateFlags = (
+        deconz_api.DeviceStateFlags.APSDE_DATA_CONFIRM
+    ),
+):
+    _, rx_schema = deconz_api.COMMAND_SCHEMAS[deconz_api.CommandId.device_state_changed]
+
+    data = deconz_api.Command(
+        command_id=deconz_api.CommandId.device_state_changed,
+        seq=api._seq,
+        payload=t.serialize_dict(
+            {
+                "status": deconz_api.Status.SUCCESS,
+                "frame_length": t.uint16_t(7),
+                "device_state": deconz_api.DeviceState(
+                    network_state=network_state,
+                    device_state=device_state,
+                ),
+                "reserved": t.uint8_t(0),
+            },
+            rx_schema,
+        ),
+    ).serialize()
+
+    asyncio.get_running_loop().call_later(0.01, api.data_received, data)
+
+
 async def test_connect(api, mock_command_rsp):
     await api.connect()
 
@@ -633,6 +662,59 @@ async def test_aps_data_request_retries_failure(api, mock_command_rsp):
     assert len(mock_rsp.mock_calls) == 1
 
 
+async def test_aps_data_request_locking(caplog, api, mock_command_rsp):
+    await api.connect()
+
+    # No free slots
+    send_network_state(api, device_state=deconz_api.DeviceStateFlags.APSDE_DATA_CONFIRM)
+
+    await asyncio.sleep(0.1)
+
+    mock_rsp = mock_command_rsp(
+        command_id=deconz_api.CommandId.aps_data_request,
+        params={},
+        rsp={
+            "status": deconz_api.Status.SUCCESS,
+            "frame_length": t.uint16_t(9),
+            "payload_length": t.uint16_t(2),
+            "device_state": deconz_api.DeviceState(
+                network_state=deconz_api.NetworkState2.CONNECTED,
+                device_state=(
+                    deconz_api.DeviceStateFlags.APSDE_DATA_REQUEST_FREE_SLOTS_AVAILABLE
+                ),
+            ),
+            "request_id": t.uint8_t(0x00),
+        },
+    )
+
+    with caplog.at_level(logging.DEBUG):
+        send = asyncio.create_task(
+            api.aps_data_request(
+                req_id=0x00,
+                dst_addr_ep=t.DeconzAddressEndpoint.deserialize(b"\x02\xaa\x55\x01")[0],
+                profile=0x0104,
+                cluster=0x0007,
+                src_ep=1,
+                aps_payload=b"aps payload",
+            )
+        )
+
+        await asyncio.sleep(0.1)
+
+    assert "Waiting for free slots to become available" in caplog.text
+
+    assert len(mock_rsp.mock_calls) == 0
+
+    send_network_state(
+        api,
+        device_state=deconz_api.DeviceStateFlags.APSDE_DATA_REQUEST_FREE_SLOTS_AVAILABLE,
+    )
+
+    await send
+
+    assert len(mock_rsp.mock_calls) == 1
+
+
 async def test_connection_lost(api):
     await api.connect()
 
@@ -787,41 +869,13 @@ async def test_data_poller(api, mock_command_rsp):
         },
     )
 
-    def send_network_state(
-        network_state: deconz_api.NetworkState2 = deconz_api.NetworkState2.CONNECTED,
-        device_state: deconz_api.DeviceStateFlags = (
-            deconz_api.DeviceStateFlags.APSDE_DATA_CONFIRM
-        ),
-    ):
-        _, rx_schema = deconz_api.COMMAND_SCHEMAS[
-            deconz_api.CommandId.device_state_changed
-        ]
-
-        data = deconz_api.Command(
-            command_id=deconz_api.CommandId.device_state_changed,
-            seq=api._seq,
-            payload=t.serialize_dict(
-                {
-                    "status": deconz_api.Status.SUCCESS,
-                    "frame_length": t.uint16_t(7),
-                    "device_state": deconz_api.DeviceState(
-                        network_state=network_state,
-                        device_state=device_state,
-                    ),
-                    "reserved": t.uint8_t(0),
-                },
-                rx_schema,
-            ),
-        ).serialize()
-
-        asyncio.get_running_loop().call_later(0.01, api.data_received, data)
-
     # Take us offline for a moment
-    send_network_state(network_state=deconz_api.NetworkState2.OFFLINE)
+    send_network_state(api, network_state=deconz_api.NetworkState2.OFFLINE)
     await asyncio.sleep(0.1)
 
     # Bring us back online with just a data confirmation to kick things off
     send_network_state(
+        api,
         network_state=deconz_api.NetworkState2.CONNECTED,
         device_state=deconz_api.DeviceStateFlags.APSDE_DATA_CONFIRM,
     )
diff --git a/zigpy_deconz/api.py b/zigpy_deconz/api.py
index 0cf8545..a6fa81c 100644
--- a/zigpy_deconz/api.py
+++ b/zigpy_deconz/api.py
@@ -702,18 +702,17 @@ def _handle_device_state_changed(
                 device_state.network_state.name,
             )
 
-        self._device_state = device_state
-        self._data_poller_event.set()
-
         if (
             DeviceStateFlags.APSDE_DATA_REQUEST_FREE_SLOTS_AVAILABLE
             in device_state.device_state
         ):
             self._free_slots_available_event.set()
         else:
-            LOGGER.debug("Data request queue full.")
             self._free_slots_available_event.clear()
 
+        self._device_state = device_state
+        self._data_poller_event.set()
+
     async def version(self):
         self._protocol_version = await self.read_parameter(
             NetworkParameter.protocol_version