Ask user for Z-Wave RF region if country is missing (#150959)

Co-authored-by: Paulus Schoutsen <balloob@gmail.com>
Co-authored-by: TheJulianJES <TheJulianJES@users.noreply.github.com>

homeassistant/components/zwave_js/config_flow.py

@@ -35,6 +35,7 @@ from homeassistant.const import CONF_NAME, CONF_URL
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.data_entry_flow import AbortFlow
 from homeassistant.exceptions import HomeAssistantError
+from homeassistant.helpers import selector
 from homeassistant.helpers.hassio import is_hassio
 from homeassistant.helpers.service_info.hassio import HassioServiceInfo
 from homeassistant.helpers.service_info.usb import UsbServiceInfo
@@ -88,6 +89,8 @@ ADDON_USER_INPUT_MAP = {
     CONF_ADDON_LR_S2_AUTHENTICATED_KEY: CONF_LR_S2_AUTHENTICATED_KEY,
 }
 
+CONF_ADDON_RF_REGION = "rf_region"
+
 EXAMPLE_SERVER_URL = "ws://localhost:3000"
 ON_SUPERVISOR_SCHEMA = vol.Schema({vol.Optional(CONF_USE_ADDON, default=True): bool})
 MIN_MIGRATION_SDK_VERSION = AwesomeVersion("6.61")
@@ -103,6 +106,19 @@ ZWAVE_JS_UI_MIGRATION_INSTRUCTIONS = (
     "#how-to-migrate-from-one-adapter-to-a-new-adapter-using-z-wave-js-ui"
 )
 
+RF_REGIONS = [
+    "Australia/New Zealand",
+    "China",
+    "Europe",
+    "Hong Kong",
+    "India",
+    "Israel",
+    "Japan",
+    "Korea",
+    "Russia",
+    "USA",
+]
+
 
 def get_manual_schema(user_input: dict[str, Any]) -> vol.Schema:
     """Return a schema for the manual step."""
@@ -195,10 +211,12 @@ class ZWaveJSConfigFlow(ConfigFlow, domain=DOMAIN):
         self.backup_data: bytes | None = None
         self.backup_filepath: Path | None = None
         self.use_addon = False
+        self._addon_config_updates: dict[str, Any] = {}
         self._migrating = False
         self._reconfigure_config_entry: ZwaveJSConfigEntry | None = None
         self._usb_discovery = False
         self._recommended_install = False
+        self._rf_region: str | None = None
 
     async def async_step_install_addon(
         self, user_input: dict[str, Any] | None = None
@@ -236,6 +254,21 @@ class ZWaveJSConfigFlow(ConfigFlow, domain=DOMAIN):
         self, user_input: dict[str, Any] | None = None
     ) -> ConfigFlowResult:
         """Start Z-Wave JS add-on."""
+        if self.hass.config.country is None and (
+            not self._rf_region or self._rf_region == "Automatic"
+        ):
+            # If the country is not set, we need to check the RF region add-on config.
+            addon_info = await self._async_get_addon_info()
+            rf_region: str | None = addon_info.options.get(CONF_ADDON_RF_REGION)
+            self._rf_region = rf_region
+            if rf_region is None or rf_region == "Automatic":
+                # If the RF region is not set, we need to ask the user to select it.
+                return await self.async_step_rf_region()
+        if config_updates := self._addon_config_updates:
+            # If we have updates to the add-on config, set them before starting the add-on.
+            self._addon_config_updates = {}
+            await self._async_set_addon_config(config_updates)
+
         if not self.start_task:
             self.start_task = self.hass.async_create_task(self._async_start_addon())
 
@@ -629,6 +662,33 @@ class ZWaveJSConfigFlow(ConfigFlow, domain=DOMAIN):
             return await self.async_step_on_supervisor({CONF_USE_ADDON: True})
         return await self.async_step_on_supervisor()
 
+    async def async_step_rf_region(
+        self, user_input: dict[str, Any] | None = None
+    ) -> ConfigFlowResult:
+        """Handle RF region selection step."""
+        if user_input is not None:
+            # Store the selected RF region
+            self._addon_config_updates[CONF_ADDON_RF_REGION] = self._rf_region = (
+                user_input["rf_region"]
+            )
+            return await self.async_step_start_addon()
+
+        schema = vol.Schema(
+            {
+                vol.Required("rf_region"): selector.SelectSelector(
+                    selector.SelectSelectorConfig(
+                        options=RF_REGIONS,
+                        mode=selector.SelectSelectorMode.DROPDOWN,
+                    )
+                ),
+            }
+        )
+
+        return self.async_show_form(
+            step_id="rf_region",
+            data_schema=schema,
+        )
+
     async def async_step_on_supervisor(
         self, user_input: dict[str, Any] | None = None
     ) -> ConfigFlowResult:
@@ -728,7 +788,7 @@ class ZWaveJSConfigFlow(ConfigFlow, domain=DOMAIN):
                     CONF_ADDON_LR_S2_AUTHENTICATED_KEY: self.lr_s2_authenticated_key,
                 }
 
-                await self._async_set_addon_config(addon_config_updates)
+                self._addon_config_updates = addon_config_updates
                 return await self.async_step_start_addon()
 
             # Network already exists, go to security keys step
@@ -799,7 +859,7 @@ class ZWaveJSConfigFlow(ConfigFlow, domain=DOMAIN):
                 CONF_ADDON_LR_S2_AUTHENTICATED_KEY: self.lr_s2_authenticated_key,
             }
 
-            await self._async_set_addon_config(addon_config_updates)
+            self._addon_config_updates = addon_config_updates
             return await self.async_step_start_addon()
 
         data_schema = vol.Schema(
@@ -1004,7 +1064,7 @@ class ZWaveJSConfigFlow(ConfigFlow, domain=DOMAIN):
         if user_input is not None:
             if self.usb_path:
                 # USB discovery was used, so the device is already known.
-                await self._async_set_addon_config({CONF_ADDON_DEVICE: self.usb_path})
+                self._addon_config_updates[CONF_ADDON_DEVICE] = self.usb_path
                 return await self.async_step_start_addon()
             # Now that the old controller is gone, we can scan for serial ports again
             return await self.async_step_choose_serial_port()
@@ -1136,6 +1196,8 @@ class ZWaveJSConfigFlow(ConfigFlow, domain=DOMAIN):
                 CONF_ADDON_LR_S2_AUTHENTICATED_KEY: self.lr_s2_authenticated_key,
             }
 
+            addon_config_updates = self._addon_config_updates | addon_config_updates
+            self._addon_config_updates = {}
             await self._async_set_addon_config(addon_config_updates)
 
             if addon_info.state == AddonState.RUNNING and not self.restart_addon:
@@ -1207,7 +1269,7 @@ class ZWaveJSConfigFlow(ConfigFlow, domain=DOMAIN):
         """Choose a serial port."""
         if user_input is not None:
             self.usb_path = user_input[CONF_USB_PATH]
-            await self._async_set_addon_config({CONF_ADDON_DEVICE: self.usb_path})
+            self._addon_config_updates[CONF_ADDON_DEVICE] = self.usb_path
             return await self.async_step_start_addon()
 
         try:

homeassistant/components/zwave_js/strings.json

@@ -113,6 +113,16 @@
         "description": "[%key:component::zwave_js::config::step::on_supervisor::description%]",
         "title": "[%key:component::zwave_js::config::step::on_supervisor::title%]"
       },
+      "rf_region": {
+        "title": "Z-Wave region",
+        "description": "Select the RF region for your Z-Wave network.",
+        "data": {
+          "rf_region": "RF region"
+        },
+        "data_description": {
+          "rf_region": "The radio frequency region for your Z-Wave network. This must match the region of your Z-Wave devices."
+        }
+      },
       "start_addon": {
         "title": "Configuring add-on"
       },

tests/components/zwave_js/test_config_flow.py

@@ -198,6 +198,17 @@ def mock_sdk_version(client: MagicMock) -> Generator[None]:
         client.driver.controller.data["sdkVersion"] = original_sdk_version
 
 
+@pytest.fixture(name="set_country", autouse=True)
+def set_country_fixture(hass: HomeAssistant) -> Generator[None]:
+    """Set the country for the test."""
+    original_country = hass.config.country
+    # Set a default country to avoid asking the user to select it.
+    hass.config.country = "US"
+    yield
+    # Reset the country after the test.
+    hass.config.country = original_country
+
+
 async def test_manual(hass: HomeAssistant) -> None:
     """Test we create an entry with manual step."""
 
@@ -4601,3 +4612,324 @@ async def test_recommended_usb_discovery(
     }
     assert len(mock_setup.mock_calls) == 1
     assert len(mock_setup_entry.mock_calls) == 1
+
+
+@pytest.mark.usefixtures("supervisor", "addon_installed", "addon_info", "unload_entry")
+async def test_addon_rf_region_new_network(
+    hass: HomeAssistant,
+    setup_entry: AsyncMock,
+    set_addon_options: AsyncMock,
+    start_addon: AsyncMock,
+) -> None:
+    """Test RF region selection for new network when country is None."""
+    device = "/test"
+    hass.config.country = None
+
+    result = await hass.config_entries.flow.async_init(
+        DOMAIN, context={"source": config_entries.SOURCE_USER}
+    )
+
+    assert result["type"] is FlowResultType.MENU
+    assert result["step_id"] == "installation_type"
+
+    result = await hass.config_entries.flow.async_configure(
+        result["flow_id"], {"next_step_id": "intent_recommended"}
+    )
+
+    result = await hass.config_entries.flow.async_configure(
+        result["flow_id"],
+        {
+            "usb_path": device,
+        },
+    )
+
+    assert result["type"] == FlowResultType.FORM
+    assert result["step_id"] == "rf_region"
+
+    # Check that all expected RF regions are available
+
+    data_schema = result["data_schema"]
+    assert data_schema is not None
+    schema = data_schema.schema
+    rf_region_field = schema["rf_region"]
+    selector_options = rf_region_field.config["options"]
+
+    expected_regions = [
+        "Australia/New Zealand",
+        "China",
+        "Europe",
+        "Hong Kong",
+        "India",
+        "Israel",
+        "Japan",
+        "Korea",
+        "Russia",
+        "USA",
+    ]
+
+    assert selector_options == expected_regions
+
+    result = await hass.config_entries.flow.async_configure(
+        result["flow_id"], {"rf_region": "Europe"}
+    )
+
+    assert result["type"] is FlowResultType.SHOW_PROGRESS
+    assert result["step_id"] == "start_addon"
+
+    # Verify RF region was set in addon config
+    assert set_addon_options.call_count == 1
+    assert set_addon_options.call_args == call(
+        "core_zwave_js",
+        AddonsOptions(
+            config={
+                "device": device,
+                "s0_legacy_key": "",
+                "s2_access_control_key": "",
+                "s2_authenticated_key": "",
+                "s2_unauthenticated_key": "",
+                "lr_s2_access_control_key": "",
+                "lr_s2_authenticated_key": "",
+                "rf_region": "Europe",
+            }
+        ),
+    )
+
+    await hass.async_block_till_done()
+    result = await hass.config_entries.flow.async_configure(result["flow_id"])
+    await hass.async_block_till_done()
+
+    assert result["type"] is FlowResultType.CREATE_ENTRY
+    assert start_addon.call_count == 1
+    assert start_addon.call_args == call("core_zwave_js")
+    assert setup_entry.call_count == 1
+
+    # avoid unload entry in teardown
+    entry = result["result"]
+    await hass.config_entries.async_unload(entry.entry_id)
+    assert entry.state is config_entries.ConfigEntryState.NOT_LOADED
+
+
+@pytest.mark.usefixtures("supervisor", "addon_running")
+async def test_addon_rf_region_migrate_network(
+    hass: HomeAssistant,
+    client: MagicMock,
+    integration: MockConfigEntry,
+    restart_addon: AsyncMock,
+    addon_options: dict[str, Any],
+    set_addon_options: AsyncMock,
+    get_server_version: AsyncMock,
+) -> None:
+    """Test migration flow with add-on."""
+    hass.config.country = None
+    version_info = get_server_version.return_value
+    entry = integration
+    assert client.connect.call_count == 1
+    assert client.driver.controller.home_id == 3245146787
+    assert entry.unique_id == "3245146787"
+    hass.config_entries.async_update_entry(
+        entry,
+        data={
+            "url": "ws://localhost:3000",
+            "use_addon": True,
+            "usb_path": "/dev/ttyUSB0",
+        },
+    )
+    addon_options["device"] = "/dev/ttyUSB0"
+
+    async def mock_backup_nvm_raw():
+        await asyncio.sleep(0)
+        client.driver.controller.emit(
+            "nvm backup progress", {"bytesRead": 100, "total": 200}
+        )
+        return b"test_nvm_data"
+
+    client.driver.controller.async_backup_nvm_raw = AsyncMock(
+        side_effect=mock_backup_nvm_raw
+    )
+
+    async def mock_restore_nvm(data: bytes, options: dict[str, bool] | None = None):
+        client.driver.controller.emit(
+            "nvm convert progress",
+            {"event": "nvm convert progress", "bytesRead": 100, "total": 200},
+        )
+        await asyncio.sleep(0)
+        client.driver.controller.emit(
+            "nvm restore progress",
+            {"event": "nvm restore progress", "bytesWritten": 100, "total": 200},
+        )
+        client.driver.controller.data["homeId"] = 3245146787
+        client.driver.emit(
+            "driver ready", {"event": "driver ready", "source": "driver"}
+        )
+
+    client.driver.controller.async_restore_nvm = AsyncMock(side_effect=mock_restore_nvm)
+
+    events = async_capture_events(
+        hass, data_entry_flow.EVENT_DATA_ENTRY_FLOW_PROGRESS_UPDATE
+    )
+
+    result = await entry.start_reconfigure_flow(hass)
+
+    assert result["type"] is FlowResultType.MENU
+    assert result["step_id"] == "reconfigure"
+
+    result = await hass.config_entries.flow.async_configure(
+        result["flow_id"], {"next_step_id": "intent_migrate"}
+    )
+
+    assert result["type"] is FlowResultType.SHOW_PROGRESS
+    assert result["step_id"] == "backup_nvm"
+
+    with patch("pathlib.Path.write_bytes") as mock_file:
+        await hass.async_block_till_done()
+        assert client.driver.controller.async_backup_nvm_raw.call_count == 1
+        assert mock_file.call_count == 1
+        assert len(events) == 1
+        assert events[0].data["progress"] == 0.5
+        events.clear()
+
+    result = await hass.config_entries.flow.async_configure(result["flow_id"])
+
+    assert result["type"] is FlowResultType.FORM
+    assert result["step_id"] == "instruct_unplug"
+    assert entry.state is config_entries.ConfigEntryState.NOT_LOADED
+
+    result = await hass.config_entries.flow.async_configure(result["flow_id"], {})
+
+    assert result["type"] is FlowResultType.FORM
+    assert result["step_id"] == "choose_serial_port"
+    data_schema = result["data_schema"]
+    assert data_schema is not None
+    assert data_schema.schema[CONF_USB_PATH]
+    # Ensure the old usb path is not in the list of options
+    with pytest.raises(InInvalid):
+        data_schema.schema[CONF_USB_PATH](addon_options["device"])
+
+    version_info.home_id = 5678
+
+    result = await hass.config_entries.flow.async_configure(
+        result["flow_id"],
+        user_input={
+            CONF_USB_PATH: "/test",
+        },
+    )
+
+    assert result["type"] == FlowResultType.FORM
+    assert result["step_id"] == "rf_region"
+
+    result = await hass.config_entries.flow.async_configure(
+        result["flow_id"], {"rf_region": "Europe"}
+    )
+
+    assert result["type"] is FlowResultType.SHOW_PROGRESS
+    assert result["step_id"] == "start_addon"
+    assert set_addon_options.call_args == call(
+        "core_zwave_js",
+        AddonsOptions(
+            config={
+                "device": "/test",
+                "rf_region": "Europe",
+            }
+        ),
+    )
+
+    await hass.async_block_till_done()
+
+    assert restart_addon.call_args == call("core_zwave_js")
+
+    result = await hass.config_entries.flow.async_configure(result["flow_id"])
+
+    assert entry.unique_id == "5678"
+    version_info.home_id = 3245146787
+
+    assert result["type"] is FlowResultType.SHOW_PROGRESS
+    assert result["step_id"] == "restore_nvm"
+    assert client.connect.call_count == 2
+
+    await hass.async_block_till_done()
+    assert client.connect.call_count == 4
+    assert entry.state is config_entries.ConfigEntryState.LOADED
+    assert client.driver.controller.async_restore_nvm.call_count == 1
+    assert len(events) == 2
+    assert events[0].data["progress"] == 0.25
+    assert events[1].data["progress"] == 0.75
+
+    result = await hass.config_entries.flow.async_configure(result["flow_id"])
+
+    assert result["type"] is FlowResultType.ABORT
+    assert result["reason"] == "migration_successful"
+    assert entry.data["url"] == "ws://host1:3001"
+    assert entry.data["usb_path"] == "/test"
+    assert entry.data["use_addon"] is True
+    assert entry.unique_id == "3245146787"
+    assert client.driver.controller.home_id == 3245146787
+
+
+@pytest.mark.usefixtures("supervisor", "addon_installed", "unload_entry")
+@pytest.mark.parametrize(("country", "rf_region"), [("US", "Automatic"), (None, "USA")])
+async def test_addon_skip_rf_region(
+    hass: HomeAssistant,
+    setup_entry: AsyncMock,
+    addon_options: dict[str, Any],
+    set_addon_options: AsyncMock,
+    start_addon: AsyncMock,
+    country: str | None,
+    rf_region: str,
+) -> None:
+    """Test RF region selection is skipped if not needed."""
+    device = "/test"
+    addon_options["rf_region"] = rf_region
+    hass.config.country = country
+
+    result = await hass.config_entries.flow.async_init(
+        DOMAIN, context={"source": config_entries.SOURCE_USER}
+    )
+
+    assert result["type"] is FlowResultType.MENU
+    assert result["step_id"] == "installation_type"
+
+    result = await hass.config_entries.flow.async_configure(
+        result["flow_id"], {"next_step_id": "intent_recommended"}
+    )
+
+    result = await hass.config_entries.flow.async_configure(
+        result["flow_id"],
+        {
+            "usb_path": device,
+        },
+    )
+
+    assert result["type"] is FlowResultType.SHOW_PROGRESS
+    assert result["step_id"] == "start_addon"
+
+    # Verify RF region was set in addon config
+    assert set_addon_options.call_count == 1
+    assert set_addon_options.call_args == call(
+        "core_zwave_js",
+        AddonsOptions(
+            config={
+                "device": device,
+                "s0_legacy_key": "",
+                "s2_access_control_key": "",
+                "s2_authenticated_key": "",
+                "s2_unauthenticated_key": "",
+                "lr_s2_access_control_key": "",
+                "lr_s2_authenticated_key": "",
+                "rf_region": rf_region,
+            }
+        ),
+    )
+
+    await hass.async_block_till_done()
+    result = await hass.config_entries.flow.async_configure(result["flow_id"])
+    await hass.async_block_till_done()
+
+    assert result["type"] is FlowResultType.CREATE_ENTRY
+    assert start_addon.call_count == 1
+    assert start_addon.call_args == call("core_zwave_js")
+    assert setup_entry.call_count == 1
+
+    # avoid unload entry in teardown
+    entry = result["result"]
+    await hass.config_entries.async_unload(entry.entry_id)
+    assert entry.state is config_entries.ConfigEntryState.NOT_LOADED