From da8a0c31fce9f94c8bbb36a377d4c7bbab9d7625 Mon Sep 17 00:00:00 2001
From: puddly <32534428+puddly@users.noreply.github.com>
Date: Fri, 9 Aug 2024 16:45:22 -0400
Subject: [PATCH] Parse EZSP `readCounters` laxly (#644)

* Parse EZSP counters laxly

* Move them up to the EZSP protocol abstraction

* Fix unit tests
---
 bellows/ezsp/protocol.py      |  8 ++++++++
 bellows/ezsp/v4/__init__.py   |  8 ++++++++
 bellows/ezsp/v4/commands.py   |  4 ++--
 bellows/zigbee/application.py |  6 +++---
 tests/test_application.py     |  6 +++---
 tests/test_ezsp_v4.py         | 16 ++++++++++++++++
 tests/test_ezsp_v7.py         |  1 +
 7 files changed, 41 insertions(+), 8 deletions(-)

diff --git a/bellows/ezsp/protocol.py b/bellows/ezsp/protocol.py
index 59367e8e..1fe87ad1 100644
--- a/bellows/ezsp/protocol.py
+++ b/bellows/ezsp/protocol.py
@@ -244,3 +244,11 @@ async def send_broadcast(
     @abc.abstractmethod
     async def set_source_route(self, nwk: t.NWK, relays: list[t.NWK]) -> t.sl_Status:
         raise NotImplementedError
+
+    @abc.abstractmethod
+    async def read_counters(self) -> dict[t.EmberCounterType, int]:
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    async def read_and_clear_counters(self) -> dict[t.EmberCounterType, int]:
+        raise NotImplementedError
diff --git a/bellows/ezsp/v4/__init__.py b/bellows/ezsp/v4/__init__.py
index fd50d269..b1e9fb6d 100644
--- a/bellows/ezsp/v4/__init__.py
+++ b/bellows/ezsp/v4/__init__.py
@@ -185,3 +185,11 @@ async def send_broadcast(
     async def set_source_route(self, nwk: t.NWK, relays: list[t.NWK]) -> t.sl_Status:
         (res,) = await self.setSourceRoute(destination=nwk, relayList=relays)
         return t.sl_Status.from_ember_status(res)
+
+    async def read_counters(self) -> dict[t.EmberCounterType, t.uint16_t]:
+        (res,) = await self.readCounters()
+        return dict(zip(t.EmberCounterType, res))
+
+    async def read_and_clear_counters(self) -> dict[t.EmberCounterType, t.uint16_t]:
+        (res,) = await self.readAndClearCounters()
+        return dict(zip(t.EmberCounterType, res))
diff --git a/bellows/ezsp/v4/commands.py b/bellows/ezsp/v4/commands.py
index 5f2fe658..64e4e4cb 100644
--- a/bellows/ezsp/v4/commands.py
+++ b/bellows/ezsp/v4/commands.py
@@ -275,14 +275,14 @@
         0x65,
         {},
         {
-            "values": t.FixedList[t.uint16_t, len(t.EmberCounterType)],
+            "values": t.List[t.uint16_t],
         },
     ),
     "readCounters": (
         0xF1,
         {},
         {
-            "values": t.FixedList[t.uint16_t, len(t.EmberCounterType)],
+            "values": t.List[t.uint16_t],
         },
     ),
     "counterRolloverHandler": (
diff --git a/bellows/zigbee/application.py b/bellows/zigbee/application.py
index b76df8fa..22247178 100644
--- a/bellows/zigbee/application.py
+++ b/bellows/zigbee/application.py
@@ -890,11 +890,11 @@ async def _watchdog_feed(self):
                 )
 
                 if remainder > 0:
-                    (res,) = await self._ezsp.readCounters()
+                    current_counters = await self._ezsp.read_counters()
                 else:
-                    (res,) = await self._ezsp.readAndClearCounters()
+                    current_counters = await self._ezsp.read_and_clear_counters()
 
-                for cnt_type, value in zip(t.EmberCounterType, res):
+                for cnt_type, value in current_counters.items():
                     counters[cnt_type.name[8:]].update(value)
 
                 if remainder == 0:
diff --git a/tests/test_application.py b/tests/test_application.py
index fe6928e4..c846f6bf 100644
--- a/tests/test_application.py
+++ b/tests/test_application.py
@@ -1210,10 +1210,10 @@ async def counters_mock():
             if nop_success % 2:
                 raise EzspError
             else:
-                return ([0, 1, 2, 3],)
+                return {t.EmberCounterType(i): v for i, v in enumerate([0, 1, 2, 3])}
         raise asyncio.TimeoutError
 
-    app._ezsp.readCounters = AsyncMock(side_effect=counters_mock)
+    app._ezsp.read_counters = AsyncMock(side_effect=counters_mock)
     app._ezsp.nop = AsyncMock(side_effect=EzspError)
     app._ezsp.getValue = AsyncMock(
         return_value=(t.EzspStatus.ERROR_OUT_OF_MEMORY, b"\x20")
@@ -1222,7 +1222,7 @@ async def counters_mock():
     app._ctrl_event.set()
 
     await app._watchdog_feed()
-    assert app._ezsp.readCounters.await_count != 0
+    assert app._ezsp.read_counters.await_count != 0
     assert app._ezsp.nop.await_count == 0
 
     cnt = t.EmberCounterType
diff --git a/tests/test_ezsp_v4.py b/tests/test_ezsp_v4.py
index d0e5fe17..69e00243 100644
--- a/tests/test_ezsp_v4.py
+++ b/tests/test_ezsp_v4.py
@@ -363,3 +363,19 @@ async def test_add_transient_link_key(ezsp_f) -> None:
         key=t.KeyData("ZigBeeAlliance09"),
     )
     assert status == t.sl_Status.OK
+
+
+@pytest.mark.parametrize("length", [40, 41])
+async def test_read_counters(ezsp_f, length: int) -> None:
+    """Test parsing of a `readCounters` response, including truncation."""
+    ezsp_f.readCounters.return_value = (list(range(length)),)
+    ezsp_f.readAndClearCounters.return_value = (list(range(length)),)
+    counters1 = await ezsp_f.read_counters()
+    counters2 = await ezsp_f.read_and_clear_counters()
+    assert (
+        ezsp_f.readCounters.mock_calls
+        == ezsp_f.readAndClearCounters.mock_calls
+        == [call()]
+    )
+
+    assert counters1 == counters2 == {t.EmberCounterType(i): i for i in range(length)}
diff --git a/tests/test_ezsp_v7.py b/tests/test_ezsp_v7.py
index 3c07f8d1..d8819c53 100644
--- a/tests/test_ezsp_v7.py
+++ b/tests/test_ezsp_v7.py
@@ -2,6 +2,7 @@
 
 import pytest
 
+from bellows.ash import DataFrame
 import bellows.ezsp.v7
 import bellows.types as t