Allow hardware integrations to specify TX power for ZHA (#155855)

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
Co-authored-by: Franck Nijhof <git@frenck.dev>

homeassistant/components/homeassistant_connect_zbt2/config_flow.py

@@ -39,6 +39,8 @@ from .const import (
     NABU_CASA_FIRMWARE_RELEASES_URL,
     PID,
     PRODUCT,
+    RADIO_TX_POWER_DBM_BY_COUNTRY,
+    RADIO_TX_POWER_DBM_DEFAULT,
     SERIAL_NUMBER,
     VID,
 )
@@ -103,6 +105,21 @@ class ZBT2FirmwareMixin(ConfigEntryBaseFlow, FirmwareInstallFlowProtocol):
             next_step_id="finish_thread_installation",
         )
 
+    def _extra_zha_hardware_options(self) -> dict[str, Any]:
+        """Return extra ZHA hardware options."""
+        country = self.hass.config.country
+
+        if country is None:
+            tx_power = RADIO_TX_POWER_DBM_DEFAULT
+        else:
+            tx_power = RADIO_TX_POWER_DBM_BY_COUNTRY.get(
+                country, RADIO_TX_POWER_DBM_DEFAULT
+            )
+
+        return {
+            "tx_power": tx_power,
+        }
+
 
 class HomeAssistantConnectZBT2ConfigFlow(
     ZBT2FirmwareMixin,

homeassistant/components/homeassistant_connect_zbt2/const.py

@@ -1,5 +1,7 @@
 """Constants for the Home Assistant Connect ZBT-2 integration."""
 
+from homeassistant.generated.countries import COUNTRIES
+
 DOMAIN = "homeassistant_connect_zbt2"
 
 NABU_CASA_FIRMWARE_RELEASES_URL = (
@@ -17,3 +19,59 @@ VID = "vid"
 DEVICE = "device"
 
 HARDWARE_NAME = "Home Assistant Connect ZBT-2"
+
+RADIO_TX_POWER_DBM_DEFAULT = 8
+RADIO_TX_POWER_DBM_BY_COUNTRY = {
+    # EU Member States
+    "AT": 10,
+    "BE": 10,
+    "BG": 10,
+    "HR": 10,
+    "CY": 10,
+    "CZ": 10,
+    "DK": 10,
+    "EE": 10,
+    "FI": 10,
+    "FR": 10,
+    "DE": 10,
+    "GR": 10,
+    "HU": 10,
+    "IE": 10,
+    "IT": 10,
+    "LV": 10,
+    "LT": 10,
+    "LU": 10,
+    "MT": 10,
+    "NL": 10,
+    "PL": 10,
+    "PT": 10,
+    "RO": 10,
+    "SK": 10,
+    "SI": 10,
+    "ES": 10,
+    "SE": 10,
+    # EEA Members
+    "IS": 10,
+    "LI": 10,
+    "NO": 10,
+    # Standards harmonized with RED or ETSI
+    "CH": 10,
+    "GB": 10,
+    "TR": 10,
+    "AL": 10,
+    "BA": 10,
+    "GE": 10,
+    "MD": 10,
+    "ME": 10,
+    "MK": 10,
+    "RS": 10,
+    "UA": 10,
+    # Other CEPT nations
+    "AD": 10,
+    "AZ": 10,
+    "MC": 10,
+    "SM": 10,
+    "VA": 10,
+}
+
+assert set(RADIO_TX_POWER_DBM_BY_COUNTRY) <= COUNTRIES

homeassistant/components/homeassistant_hardware/firmware_config_flow.py

@@ -456,6 +456,10 @@ class BaseFirmwareInstallFlow(ConfigEntryBaseFlow, ABC):
         # This step is necessary to prevent `user_input` from being passed through
         return await self.async_step_continue_zigbee()
 
+    def _extra_zha_hardware_options(self) -> dict[str, Any]:
+        """Return extra ZHA hardware options."""
+        return {}
+
     async def async_step_continue_zigbee(
         self, user_input: dict[str, Any] | None = None
     ) -> ConfigFlowResult:
@@ -478,6 +482,7 @@ class BaseFirmwareInstallFlow(ConfigEntryBaseFlow, ABC):
                 },
                 "radio_type": "ezsp",
                 "flow_strategy": self._zigbee_flow_strategy,
+                **self._extra_zha_hardware_options(),
             },
         )
         return self._continue_zha_flow(result)

homeassistant/components/zha/config_flow.py

@@ -14,7 +14,7 @@ from typing import Any
 import voluptuous as vol
 from zha.application.const import RadioType
 import zigpy.backups
-from zigpy.config import CONF_DEVICE, CONF_DEVICE_PATH
+from zigpy.config import CONF_DEVICE, CONF_DEVICE_PATH, CONF_NWK_TX_POWER
 from zigpy.exceptions import CannotWriteNetworkSettings, DestructiveWriteNetworkSettings
 
 from homeassistant.components import onboarding, usb
@@ -191,6 +191,7 @@ class BaseZhaFlow(ConfigEntryBaseFlow):
         self._hass = None  # type: ignore[assignment]
         self._radio_mgr = ZhaRadioManager()
         self._restore_backup_task: asyncio.Task[None] | None = None
+        self._extra_network_config: dict[str, Any] = {}
 
     @property
     def hass(self) -> HomeAssistant:
@@ -622,7 +623,8 @@ class BaseZhaFlow(ConfigEntryBaseFlow):
         self, user_input: dict[str, Any] | None = None
     ) -> ConfigFlowResult:
         """Form a brand-new network."""
-        await self._radio_mgr.async_form_network()
+        await self._radio_mgr.async_form_network(config=self._extra_network_config)
+
         # Load the newly formed network settings to get the network info
         await self._radio_mgr.async_load_network_settings()
         return await self._async_create_radio_entry()
@@ -1007,6 +1009,9 @@ class ZhaConfigFlowHandler(BaseZhaFlow, ConfigFlow, domain=DOMAIN):
         device_path = device_settings[CONF_DEVICE_PATH]
         self._flow_strategy = discovery_data.get("flow_strategy")
 
+        if "tx_power" in discovery_data:
+            self._extra_network_config[CONF_NWK_TX_POWER] = discovery_data["tx_power"]
+
         await self._set_unique_id_and_update_ignored_flow(
             unique_id=f"{name}_{radio_type.name}_{device_path}",
             device_path=device_path,

homeassistant/components/zha/radio_manager.py

@@ -78,6 +78,7 @@ HARDWARE_DISCOVERY_SCHEMA = vol.Schema(
         vol.Required("port"): DEVICE_SCHEMA,
         vol.Required("radio_type"): str,
         vol.Optional("flow_strategy"): vol.All(str, vol.Coerce(ZigbeeFlowStrategy)),
+        vol.Optional("tx_power"): vol.All(vol.Coerce(int), vol.Range(min=0, max=10)),
     }
 )
 
@@ -322,7 +323,7 @@ class ZhaRadioManager:
 
         return backup
 
-    async def async_form_network(self) -> None:
+    async def async_form_network(self, config: dict[str, Any] | None) -> None:
         """Form a brand-new network."""
 
         # When forming a new network, we delete the ZHA database to prevent old devices
@@ -331,7 +332,7 @@ class ZhaRadioManager:
             await self.hass.async_add_executor_job(os.remove, self.zigpy_database_path)
 
         async with self.create_zigpy_app() as app:
-            await app.form_network()
+            await app.form_network(config=config)
 
     async def async_reset_adapter(self) -> None:
         """Reset the current adapter."""

tests/components/homeassistant_connect_zbt2/test_config_flow.py

@@ -49,10 +49,23 @@ def setup_entry_fixture() -> Generator[AsyncMock]:
         yield mock_setup_entry
 
 
+@pytest.mark.parametrize(
+    ("country", "expected_tx_power"),
+    [
+        ("US", 8),
+        ("NL", 10),
+        ("JP", 8),
+        ("DE", 10),
+    ],
+)
 async def test_config_flow_zigbee(
     hass: HomeAssistant,
+    country: str,
+    expected_tx_power: int,
 ) -> None:
     """Test Zigbee config flow for Connect ZBT-2."""
+    hass.config.country = country
+
     fw_type = ApplicationType.EZSP
     fw_version = "7.4.4.0 build 0"
     model = "Home Assistant Connect ZBT-2"
@@ -146,6 +159,7 @@ async def test_config_flow_zigbee(
             "flow_control": "hardware",
         },
         "radio_type": fw_type.value,
+        "tx_power": expected_tx_power,
     }
 
 
@@ -382,6 +396,7 @@ async def test_options_flow(
             "flow_control": "hardware",
         },
         "radio_type": "ezsp",
+        "tx_power": 8,
     }