From ec398b9958938dba29fc57943f2cb30ebdc6c9d2 Mon Sep 17 00:00:00 2001
From: puddly <32534428+puddly@users.noreply.github.com>
Date: Wed, 14 Aug 2024 12:55:09 -0400
Subject: [PATCH] Cache the address table size once it's read

---
 bellows/ezsp/protocol.py    |  3 +++
 bellows/ezsp/v4/__init__.py | 23 +++++++++++++----------
 tests/test_ezsp_v4.py       | 29 +++++++++++++++++++++++++++++
 3 files changed, 45 insertions(+), 10 deletions(-)

diff --git a/bellows/ezsp/protocol.py b/bellows/ezsp/protocol.py
index 0e6ad367..f9eca74e 100644
--- a/bellows/ezsp/protocol.py
+++ b/bellows/ezsp/protocol.py
@@ -43,6 +43,9 @@ def __init__(self, cb_handler: Callable, gateway: Gateway) -> None:
         }
         self.tc_policy = 0
 
+        # Cached by `set_extended_timeout` so subsequent calls are a little faster
+        self._address_table_size: int | None = None
+
     def _ezsp_frame(self, name: str, *args: Any, **kwargs: Any) -> bytes:
         """Serialize the named frame and data."""
         c, tx_schema, rx_schema = self.COMMANDS[name]
diff --git a/bellows/ezsp/v4/__init__.py b/bellows/ezsp/v4/__init__.py
index ebf18001..534c842a 100644
--- a/bellows/ezsp/v4/__init__.py
+++ b/bellows/ezsp/v4/__init__.py
@@ -212,19 +212,22 @@ async def set_extended_timeout(
             )
             return
 
-        (status, addr_table_size) = await self.getConfigurationValue(
-            t.EzspConfigId.CONFIG_ADDRESS_TABLE_SIZE
-        )
-
-        if t.sl_Status.from_ember_status(status) != t.sl_Status.OK:
-            # Last-ditch effort
-            await self.setExtendedTimeout(
-                remoteEui64=ieee, extendedTimeout=extended_timeout
+        if self._address_table_size is None:
+            (status, addr_table_size) = await self.getConfigurationValue(
+                t.EzspConfigId.CONFIG_ADDRESS_TABLE_SIZE
             )
-            return
+
+            if t.sl_Status.from_ember_status(status) != t.sl_Status.OK:
+                # Last-ditch effort
+                await self.setExtendedTimeout(
+                    remoteEui64=ieee, extendedTimeout=extended_timeout
+                )
+                return
+
+            self._address_table_size = addr_table_size
 
         # Replace a random entry in the address table
-        index = random.randint(0, addr_table_size - 1)
+        index = random.randint(0, self._address_table_size - 1)
 
         await self.replaceAddressTableEntry(
             addressTableIndex=index,
diff --git a/tests/test_ezsp_v4.py b/tests/test_ezsp_v4.py
index 3f035366..1fec0187 100644
--- a/tests/test_ezsp_v4.py
+++ b/tests/test_ezsp_v4.py
@@ -420,6 +420,35 @@ async def test_set_extended_timeout_no_entry(ezsp_f) -> None:
         )
     ]
 
+    # The address table size is cached
+    with patch("bellows.ezsp.v4.random.randint") as mock_random:
+        mock_random.return_value = 1
+        await ezsp_f.set_extended_timeout(
+            nwk=0x1234,
+            ieee=t.EUI64.convert("aa:bb:cc:dd:ee:ff:00:11"),
+            extended_timeout=True,
+        )
+
+    # Still called only once
+    assert ezsp_f.getConfigurationValue.mock_calls == [
+        call(t.EzspConfigId.CONFIG_ADDRESS_TABLE_SIZE)
+    ]
+
+    assert ezsp_f.replaceAddressTableEntry.mock_calls == [
+        call(
+            addressTableIndex=0,
+            newEui64=t.EUI64.convert("aa:bb:cc:dd:ee:ff:00:11"),
+            newId=0x1234,
+            newExtendedTimeout=True,
+        ),
+        call(
+            addressTableIndex=1,
+            newEui64=t.EUI64.convert("aa:bb:cc:dd:ee:ff:00:11"),
+            newId=0x1234,
+            newExtendedTimeout=True,
+        ),
+    ]
+
 
 async def test_set_extended_timeout_already_set(ezsp_f) -> None:
     # No-op, it's already set