Implement zigpy TX power API (#698)

* Implement zigpy TX power API

* Fix typo in command name

* Bump zigpy

* Add tests

Author: puddly

bellows/ezsp/__init__.py

@@ -20,7 +20,12 @@
 from bellows.exception import EzspError, InvalidCommandError, InvalidCommandPayload
 from bellows.ezsp import xncp
 from bellows.ezsp.config import DEFAULT_CONFIG, RuntimeConfig, ValueConfig
-from bellows.ezsp.xncp import FirmwareFeatures, FlowControlType, GetRouteTableEntryRsp
+from bellows.ezsp.xncp import (
+    FirmwareFeatures,
+    FlowControlType,
+    GetRouteTableEntryRsp,
+    GetTxPowerInfoRsp,
+)
 import bellows.types as t
 import bellows.uart
 
@@ -842,3 +847,8 @@ async def xncp_set_route_table_entry(
                 cost=cost,
             )
         )
+
+    async def xncp_get_tx_power_info(self, country_code: str) -> GetTxPowerInfoRsp:
+        """Get maximum and recommended TX power for a country (ISO 3166-1 alpha-2)."""
+        code = country_code.upper().encode("ascii")
+        return await self.send_xncp_frame(xncp.GetTxPowerInfoReq(country_code=code))

bellows/ezsp/xncp.py

@@ -42,6 +42,7 @@ class XncpCommandId(t.enum16):
     GET_CHIP_INFO_REQ = 0x0005
     SET_ROUTE_TABLE_ENTRY_REQ = 0x0006
     GET_ROUTE_TABLE_ENTRY_REQ = 0x0007
+    GET_TX_POWER_INFO_REQ = 0x0008
 
     GET_SUPPORTED_FEATURES_RSP = GET_SUPPORTED_FEATURES_REQ | 0x8000
     SET_SOURCE_ROUTE_RSP = SET_SOURCE_ROUTE_REQ | 0x8000
@@ -51,6 +52,7 @@ class XncpCommandId(t.enum16):
     GET_CHIP_INFO_RSP = GET_CHIP_INFO_REQ | 0x8000
     SET_ROUTE_TABLE_ENTRY_RSP = SET_ROUTE_TABLE_ENTRY_REQ | 0x8000
     GET_ROUTE_TABLE_ENTRY_RSP = GET_ROUTE_TABLE_ENTRY_REQ | 0x8000
+    GET_TX_POWER_INFO_RSP = GET_TX_POWER_INFO_REQ | 0x8000
 
     UNKNOWN = 0xFFFF
 
@@ -118,6 +120,9 @@ class FirmwareFeatures(t.bitmap32):
     # Route table entries can be set
     RESTORE_ROUTE_TABLE = 1 << 6
 
+    # Recommended and maximum TX power can be queried by country code
+    TX_POWER_INFO = 1 << 7
+
 
 class XncpCommandPayload(t.Struct):
     pass
@@ -217,6 +222,17 @@ class GetRouteTableEntryRsp(XncpCommandPayload):
     cost: t.uint8_t
 
 
+@register_command(XncpCommandId.GET_TX_POWER_INFO_REQ)
+class GetTxPowerInfoReq(XncpCommandPayload):
+    country_code: Bytes
+
+
+@register_command(XncpCommandId.GET_TX_POWER_INFO_RSP)
+class GetTxPowerInfoRsp(XncpCommandPayload):
+    recommended_power_dbm: t.int8s
+    max_power_dbm: t.int8s
+
+
 @register_command(XncpCommandId.UNKNOWN)
 class Unknown(XncpCommandPayload):
     pass

bellows/types/struct.py

@@ -376,6 +376,18 @@ class NV3StackTrustCenterToken(EzspStruct):
     key: named.KeyData
 
 
+class NV3StackNodeData(EzspStruct):
+    """NV3 stack node data token value."""
+
+    panId: named.EmberPanId
+    radioTxPower: basic.int8s
+    radioFreqChannel: basic.uint8_t
+    stackProfile: basic.uint8_t  # Always 0x02
+    nodeType: named.EmberNodeType
+    zigbeeNodeId: named.EmberNodeId
+    extendedPanId: named.ExtendedPanId
+
+
 class EmberKeyStruct(EzspStruct):
     # A structure containing a key and its associated data.
     # A bitmask indicating the presence of data within the various fields

bellows/zigbee/application.py

@@ -153,6 +153,34 @@ async def _get_board_info(self) -> tuple[str, str, str] | tuple[None, None, None
 
         return None, None, None
 
+    async def _get_recommended_tx_power(self, country: str) -> float:
+        """Get firmware recommended TX power for the given country."""
+        if FirmwareFeatures.TX_POWER_INFO not in self._ezsp._xncp_features:
+            return await super()._get_recommended_tx_power(country)
+
+        tx_power_info = await self._ezsp.xncp_get_tx_power_info(country)
+        return tx_power_info.recommended_power_dbm
+
+    async def _get_maximum_tx_power(self, country: str) -> float:
+        """Get firmware maximum TX power for the given country."""
+        if FirmwareFeatures.TX_POWER_INFO not in self._ezsp._xncp_features:
+            return await super()._get_maximum_tx_power(country)
+
+        tx_power_info = await self._ezsp.xncp_get_tx_power_info(country)
+        return tx_power_info.max_power_dbm
+
+    async def _set_tx_power(self, tx_power: float) -> float | None:
+        """Set TX power (if supported by the radio), returning the actual TX power."""
+        actual_power = int(tx_power)
+        await self._ezsp.setRadioPower(power=actual_power)
+
+        # We intentionally do not reset after changing the TX power. Instead, we just
+        # persist the changes to NVRAM (if necessary), they will be reloaded on next
+        # boot.
+        await repairs.update_tx_power(self._ezsp, tx_power=actual_power)
+
+        return float(actual_power)
+
     async def connect(self) -> None:
         self._ezsp = bellows.ezsp.EZSP(self.config[zigpy.config.CONF_DEVICE], self)
 

bellows/zigbee/repairs.py

@@ -49,3 +49,42 @@ async def fix_invalid_tclk_partner_ieee(ezsp: EZSP) -> bool:
     assert t.sl_Status.from_ember_status(status) == t.sl_Status.OK
 
     return True
+
+
+async def update_tx_power(ezsp: EZSP, tx_power: int) -> bool:
+    """Persist transmit power in NVRAM."""
+
+    try:
+        rsp = await ezsp.getTokenData(token=t.NV3KeyId.NVM3KEY_STACK_NODE_DATA, index=0)
+        assert t.sl_Status.from_ember_status(rsp.status) == t.sl_Status.OK
+    except (InvalidCommandError, AttributeError, AssertionError):
+        LOGGER.warning("NV3 interface not available in this firmware, please upgrade!")
+        return False
+
+    token, remaining = t.NV3StackNodeData.deserialize(rsp.value)
+    assert not remaining
+
+    # No point in writing to NVRAM if the TX power is correct
+    if token.radioTxPower == tx_power:
+        return False
+
+    status, node_type, nwk_params = await ezsp.getNetworkParameters()
+    assert t.sl_Status.from_ember_status(status) == t.sl_Status.OK
+
+    # Sanity check
+    assert token.panId == nwk_params.panId
+    assert token.radioFreqChannel == nwk_params.radioChannel
+    assert token.stackProfile == 0x02
+    assert token.nodeType == node_type
+    assert token.extendedPanId == nwk_params.extendedPanId
+
+    (status,) = await ezsp.setTokenData(
+        token=t.NV3KeyId.NVM3KEY_STACK_NODE_DATA,
+        index=0,
+        token_data=token.replace(radioTxPower=tx_power).serialize(),
+    )
+    assert t.sl_Status.from_ember_status(status) == t.sl_Status.OK
+
+    LOGGER.debug("Persisted TX power %d to NVRAM", tx_power)
+
+    return True

pyproject.toml

@@ -17,7 +17,7 @@ dependencies = [
     "click",
     "click-log>=0.2.1",
     "voluptuous",
-    "zigpy>=0.85.0",
+    "zigpy>=0.87.0",
 ]
 
 [tool.setuptools.packages.find]

tests/test_application.py

@@ -22,6 +22,7 @@
     FlowControlType,
     GetChipInfoRsp,
     GetRouteTableEntryRsp,
+    GetTxPowerInfoRsp,
 )
 import bellows.types
 import bellows.types as t
@@ -2639,3 +2640,46 @@ async def test_migration_failure_eui64_overwrite_confirmation(
         assert app._ezsp.write_custom_eui64.mock_calls == [
             call(t.EUI64.convert("aa:aa:aa:aa:aa:aa:aa:aa"), burn_into_userdata=True)
         ]
+
+
+async def test_tx_power_with_xncp_feature(app: ControllerApplication) -> None:
+    """Test TX power methods with XNCP TX_POWER_INFO feature."""
+    app._ezsp._xncp_features |= FirmwareFeatures.TX_POWER_INFO
+    app._ezsp.xncp_get_tx_power_info = AsyncMock(
+        return_value=GetTxPowerInfoRsp(recommended_power_dbm=10, max_power_dbm=20)
+    )
+
+    assert await app.get_recommended_tx_power("US") == 10.0
+    assert await app.get_maximum_tx_power("US") == 20.0
+
+
+async def test_tx_power_without_xncp_feature(app: ControllerApplication) -> None:
+    """Test TX power methods fall back to parent class without XNCP feature."""
+    app._ezsp._xncp_features = FirmwareFeatures.NONE
+    app._ezsp.xncp_get_tx_power_info = AsyncMock()
+
+    app_cls = zigpy.application.ControllerApplication
+
+    ezsp_rec_tx_power = await app.get_recommended_tx_power("US")
+    base_rec_tx_power = await app_cls.get_recommended_tx_power(app, "US")
+    assert ezsp_rec_tx_power == base_rec_tx_power
+
+    ezsp_max_tx_power = await app.get_maximum_tx_power("US")
+    base_max_tx_power = await app_cls.get_maximum_tx_power(app, "US")
+    assert ezsp_max_tx_power == base_max_tx_power
+
+    assert len(app._ezsp.xncp_get_tx_power_info.mock_calls) == 0
+
+
+async def test_set_tx_power(app: ControllerApplication) -> None:
+    """Test set_tx_power with float-to-int conversion and NVRAM persistence."""
+    app._ezsp.setRadioPower = AsyncMock()
+
+    with patch(
+        "bellows.zigbee.repairs.update_tx_power", return_value=True
+    ) as mock_update:
+        result = await app.set_tx_power(12.7)
+
+    assert result == 12.0
+    assert app._ezsp.setRadioPower.mock_calls == [call(power=12)]
+    assert mock_update.mock_calls == [call(app._ezsp, tx_power=12)]

tests/test_xncp.py

@@ -291,3 +291,26 @@ async def test_xncp_route_table_operations(ezsp_f: EZSP) -> None:
             ).serialize()
         )
     ]
+
+
+async def test_xncp_get_tx_power_info(ezsp_f: EZSP) -> None:
+    """Test XNCP get_tx_power_info."""
+    ezsp_f._mock_commands["customFrame"] = customFrame = AsyncMock(
+        return_value=[
+            t.EmberStatus.SUCCESS,
+            xncp.XncpCommand.from_payload(
+                xncp.GetTxPowerInfoRsp(recommended_power_dbm=10, max_power_dbm=20)
+            ).serialize(),
+        ]
+    )
+
+    rsp = await ezsp_f.xncp_get_tx_power_info("us")
+    assert rsp.recommended_power_dbm == 10
+    assert rsp.max_power_dbm == 20
+    assert customFrame.mock_calls == [
+        call(
+            xncp.XncpCommand.from_payload(
+                xncp.GetTxPowerInfoReq(country_code=b"US")
+            ).serialize()
+        )
+    ]

tests/test_zigbee_repairs.py

@@ -174,3 +174,69 @@ async def test_fix_invalid_tclk_all_versions(
         ]
     else:
         assert "NV3 interface not available in this firmware" in caplog.text
+
+
+async def test_update_tx_power(ezsp_f: EZSP, caplog) -> None:
+    """Test update_tx_power behavior in various scenarios."""
+    token_data = t.NV3StackNodeData(
+        panId=t.EmberPanId(0x1234),
+        radioTxPower=t.int8s(5),
+        radioFreqChannel=t.uint8_t(15),
+        stackProfile=t.uint8_t(0x02),
+        nodeType=t.EmberNodeType.COORDINATOR,
+        zigbeeNodeId=t.EmberNodeId(0x0000),
+        extendedPanId=t.ExtendedPanId.convert("AA:BB:CC:DD:EE:FF:00:11"),
+    )
+
+    # Test 1: NV3 interface unavailable
+    ezsp_f.getTokenData = AsyncMock(side_effect=InvalidCommandError())
+    with caplog.at_level(logging.WARNING):
+        assert await repairs.update_tx_power(ezsp_f, tx_power=10) is False
+    assert "NV3 interface not available in this firmware" in caplog.text
+
+    # Test 2: TX power already correct (no write needed)
+    ezsp_f.getTokenData = AsyncMock(
+        return_value=GetTokenDataRsp(
+            status=t.EmberStatus.SUCCESS,
+            value=token_data.replace(radioTxPower=t.int8s(10)).serialize(),
+        )
+    )
+    ezsp_f.setTokenData = AsyncMock()
+    ezsp_f.getNetworkParameters = AsyncMock()
+    assert await repairs.update_tx_power(ezsp_f, tx_power=10) is False
+    assert len(ezsp_f.setTokenData.mock_calls) == 0
+    assert len(ezsp_f.getNetworkParameters.mock_calls) == 0
+
+    # Test 3: Successful TX power update
+    ezsp_f.getTokenData = AsyncMock(
+        return_value=GetTokenDataRsp(
+            status=t.EmberStatus.SUCCESS,
+            value=token_data.serialize(),
+        )
+    )
+    ezsp_f.getNetworkParameters = AsyncMock(
+        return_value=[
+            t.EmberStatus.SUCCESS,
+            t.EmberNodeType.COORDINATOR,
+            t.EmberNetworkParameters(
+                panId=t.EmberPanId(0x1234),
+                extendedPanId=t.ExtendedPanId.convert("AA:BB:CC:DD:EE:FF:00:11"),
+                radioChannel=t.uint8_t(15),
+                radioTxPower=t.int8s(5),
+                joinMethod=t.EmberJoinMethod.USE_MAC_ASSOCIATION,
+                nwkManagerId=t.EmberNodeId(0x0000),
+                nwkUpdateId=t.uint8_t(0),
+                channels=t.Channels.ALL_CHANNELS,
+            ),
+        ]
+    )
+    ezsp_f.setTokenData = AsyncMock(return_value=[t.EmberStatus.SUCCESS])
+
+    assert await repairs.update_tx_power(ezsp_f, tx_power=15) is True
+    assert ezsp_f.setTokenData.mock_calls == [
+        call(
+            token=t.NV3KeyId.NVM3KEY_STACK_NODE_DATA,
+            index=0,
+            token_data=token_data.replace(radioTxPower=t.int8s(15)).serialize(),
+        )
+    ]