Add config entry migration for Z-Wave JS (#179898)

Co-authored-by: Paulus Schoutsen <balloob@gmail.com>
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>

homeassistant/components/zwave_js/__init__.py

@@ -162,17 +162,31 @@ PLATFORMS = [
 
 async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
     """Set up the Z-Wave JS component."""
-    for entry in hass.config_entries.async_entries(DOMAIN):
-        if not isinstance(entry.unique_id, str):
-            hass.config_entries.async_update_entry(
-                entry, unique_id=str(entry.unique_id)
-            )
-
     async_setup_services(hass)
 
     return True
 
 
+async def async_migrate_entry(hass: HomeAssistant, entry: ZwaveJSConfigEntry) -> bool:
+    """Migrate old config entry."""
+    if entry.version == 1 and entry.minor_version < 2:
+        unique_id = entry.unique_id
+        if not isinstance(unique_id, str):
+            # Old entries stored the home ID as int.
+            unique_id = str(unique_id)
+        data = dict(entry.data)
+        # s0_legacy_key was saved as network_key before s2 was added.
+        if CONF_NETWORK_KEY in data:
+            network_key = data.pop(CONF_NETWORK_KEY)
+            if not data.get(CONF_S0_LEGACY_KEY):
+                data[CONF_S0_LEGACY_KEY] = network_key
+        hass.config_entries.async_update_entry(
+            entry, data=data, unique_id=unique_id, minor_version=2
+        )
+
+    return True
+
+
 async def async_setup_entry(hass: HomeAssistant, entry: ZwaveJSConfigEntry) -> bool:
     """Set up Z-Wave JS from a config entry."""
     if use_addon := entry.data.get(CONF_USE_ADDON):
@@ -1210,10 +1224,7 @@ async def async_ensure_addon_running(
 
     usb_path: str | None = entry.data[CONF_USB_PATH]
     socket_path: str | None = entry.data.get(CONF_SOCKET_PATH)
-    # s0_legacy_key was saved as network_key before s2 was added.
     s0_legacy_key: str = entry.data.get(CONF_S0_LEGACY_KEY, "")
-    if not s0_legacy_key:
-        s0_legacy_key = entry.data.get(CONF_NETWORK_KEY, "")
     s2_access_control_key: str = entry.data.get(CONF_S2_ACCESS_CONTROL_KEY, "")
     s2_authenticated_key: str = entry.data.get(CONF_S2_AUTHENTICATED_KEY, "")
     s2_unauthenticated_key: str = entry.data.get(CONF_S2_UNAUTHENTICATED_KEY, "")

homeassistant/components/zwave_js/config_flow.py

@@ -390,6 +390,7 @@ class ZWaveJSConfigFlow(ConfigFlow, domain=DOMAIN):
         return AddonFlowManager(self.hass)
 
     VERSION = 1
+    MINOR_VERSION = 2
 
     def __init__(self) -> None:
         """Set up flow instance."""

tests/components/zwave_js/test_init.py

@@ -62,6 +62,74 @@ def connect_timeout_fixture() -> Generator[int]:
         yield timeout
 
 
+@pytest.mark.parametrize(
+    ("unique_id", "data", "expected_unique_id", "expected_data"),
+    [
+        pytest.param(
+            3245146787,
+            {"url": "ws://test.org"},
+            "3245146787",
+            {"url": "ws://test.org"},
+            id="int_unique_id",
+        ),
+        pytest.param(
+            "3245146787",
+            {"url": "ws://test.org", "network_key": "abc123"},
+            "3245146787",
+            {"url": "ws://test.org", "s0_legacy_key": "abc123"},
+            id="network_key_only",
+        ),
+        pytest.param(
+            "3245146787",
+            {
+                "url": "ws://test.org",
+                "network_key": "abc123",
+                "s0_legacy_key": "def456",
+            },
+            "3245146787",
+            {"url": "ws://test.org", "s0_legacy_key": "def456"},
+            id="existing_s0_legacy_key_wins",
+        ),
+    ],
+)
+@pytest.mark.usefixtures("client")
+async def test_migrate_entry(
+    hass: HomeAssistant,
+    unique_id: int | str,
+    data: dict[str, Any],
+    expected_unique_id: str,
+    expected_data: dict[str, Any],
+) -> None:
+    """Test migration of a version 1.1 config entry."""
+    entry = MockConfigEntry(
+        domain=DOMAIN, data=data, unique_id=unique_id, minor_version=1
+    )
+    entry.add_to_hass(hass)
+
+    with patch("homeassistant.components.zwave_js.PLATFORMS", []):
+        await hass.config_entries.async_setup(entry.entry_id)
+        await hass.async_block_till_done()
+
+    assert entry.state is ConfigEntryState.LOADED
+    assert entry.version == 1
+    assert entry.minor_version == 2
+    assert entry.unique_id == expected_unique_id
+    assert dict(entry.data) == expected_data
+
+
+async def test_migrate_entry_from_future_version(hass: HomeAssistant) -> None:
+    """Test migration of a config entry from a future version fails."""
+    entry = MockConfigEntry(
+        domain=DOMAIN, data={"url": "ws://test.org"}, unique_id="3245146787", version=2
+    )
+    entry.add_to_hass(hass)
+
+    await hass.config_entries.async_setup(entry.entry_id)
+    await hass.async_block_till_done()
+
+    assert entry.state is ConfigEntryState.MIGRATION_ERROR
+
+
 async def test_entry_setup_unload(
     hass: HomeAssistant,
     client: MagicMock,