Cache the address table size once it's read

zigpy · Aug 14, 2024 · ec398b9 · ec398b9
1 parent 2d4605a
commit ec398b9
Show file tree

Hide file tree

Showing 3 changed files with 45 additions and 10 deletions.
diff --git 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
@@ -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
@@ -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