Fix Z-Wave removal of devices when connected to unknown controller (#149339)

homeassistant/components/zwave_js/__init__.py

@@ -277,39 +277,6 @@ async def async_setup_entry(hass: HomeAssistant, entry: ZwaveJSConfigEntry) -> b
         # and we'll handle the clean up below.
         await driver_events.setup(driver)
 
-    if (old_unique_id := entry.unique_id) is not None and old_unique_id != (
-        new_unique_id := str(driver.controller.home_id)
-    ):
-        device_registry = dr.async_get(hass)
-        controller_model = "Unknown model"
-        if (
-            (own_node := driver.controller.own_node)
-            and (
-                controller_device_entry := device_registry.async_get_device(
-                    identifiers={get_device_id(driver, own_node)}
-                )
-            )
-            and (model := controller_device_entry.model)
-        ):
-            controller_model = model
-        async_create_issue(
-            hass,
-            DOMAIN,
-            f"migrate_unique_id.{entry.entry_id}",
-            data={
-                "config_entry_id": entry.entry_id,
-                "config_entry_title": entry.title,
-                "controller_model": controller_model,
-                "new_unique_id": new_unique_id,
-                "old_unique_id": old_unique_id,
-            },
-            is_fixable=True,
-            severity=IssueSeverity.ERROR,
-            translation_key="migrate_unique_id",
-        )
-    else:
-        async_delete_issue(hass, DOMAIN, f"migrate_unique_id.{entry.entry_id}")
-
     # If the listen task is already failed, we need to raise ConfigEntryNotReady
     if listen_task.done():
         listen_error, error_message = _get_listen_task_error(listen_task)
@@ -387,28 +354,6 @@ class DriverEvents:
             self.hass.bus.async_listen(EVENT_LOGGING_CHANGED, handle_logging_changed)
         )
 
-        # Check for nodes that no longer exist and remove them
-        stored_devices = dr.async_entries_for_config_entry(
-            self.dev_reg, self.config_entry.entry_id
-        )
-        known_devices = [
-            self.dev_reg.async_get_device(identifiers={get_device_id(driver, node)})
-            for node in controller.nodes.values()
-        ]
-        provisioned_devices = [
-            self.dev_reg.async_get(entry.additional_properties["device_id"])
-            for entry in await controller.async_get_provisioning_entries()
-            if entry.additional_properties
-            and "device_id" in entry.additional_properties
-        ]
-
-        # Devices that are in the device registry that are not known by the controller
-        # can be removed
-        if not self.config_entry.data.get(CONF_KEEP_OLD_DEVICES):
-            for device in stored_devices:
-                if device not in known_devices and device not in provisioned_devices:
-                    self.dev_reg.async_remove_device(device.id)
-
         # run discovery on controller node
         if controller.own_node:
             await self.controller_events.async_on_node_added(controller.own_node)
@@ -443,6 +388,72 @@ class DriverEvents:
             controller.on("identify", self.controller_events.async_on_identify)
         )
 
+        if (
+            old_unique_id := self.config_entry.unique_id
+        ) is not None and old_unique_id != (
+            new_unique_id := str(driver.controller.home_id)
+        ):
+            device_registry = dr.async_get(self.hass)
+            controller_model = "Unknown model"
+            if (
+                (own_node := driver.controller.own_node)
+                and (
+                    controller_device_entry := device_registry.async_get_device(
+                        identifiers={get_device_id(driver, own_node)}
+                    )
+                )
+                and (model := controller_device_entry.model)
+            ):
+                controller_model = model
+
+            # Do not clean up old stale devices if an unknown controller is connected.
+            data = {**self.config_entry.data, CONF_KEEP_OLD_DEVICES: True}
+            self.hass.config_entries.async_update_entry(self.config_entry, data=data)
+            async_create_issue(
+                self.hass,
+                DOMAIN,
+                f"migrate_unique_id.{self.config_entry.entry_id}",
+                data={
+                    "config_entry_id": self.config_entry.entry_id,
+                    "config_entry_title": self.config_entry.title,
+                    "controller_model": controller_model,
+                    "new_unique_id": new_unique_id,
+                    "old_unique_id": old_unique_id,
+                },
+                is_fixable=True,
+                severity=IssueSeverity.ERROR,
+                translation_key="migrate_unique_id",
+            )
+        else:
+            data = self.config_entry.data.copy()
+            data.pop(CONF_KEEP_OLD_DEVICES, None)
+            self.hass.config_entries.async_update_entry(self.config_entry, data=data)
+            async_delete_issue(
+                self.hass, DOMAIN, f"migrate_unique_id.{self.config_entry.entry_id}"
+            )
+
+        # Check for nodes that no longer exist and remove them
+        stored_devices = dr.async_entries_for_config_entry(
+            self.dev_reg, self.config_entry.entry_id
+        )
+        known_devices = [
+            self.dev_reg.async_get_device(identifiers={get_device_id(driver, node)})
+            for node in controller.nodes.values()
+        ]
+        provisioned_devices = [
+            self.dev_reg.async_get(entry.additional_properties["device_id"])
+            for entry in await controller.async_get_provisioning_entries()
+            if entry.additional_properties
+            and "device_id" in entry.additional_properties
+        ]
+
+        # Devices that are in the device registry that are not known by the controller
+        # can be removed
+        if not self.config_entry.data.get(CONF_KEEP_OLD_DEVICES):
+            for device in stored_devices:
+                if device not in known_devices and device not in provisioned_devices:
+                    self.dev_reg.async_remove_device(device.id)
+
 
 class ControllerEvents:
     """Represent controller events.

homeassistant/components/zwave_js/repairs.py

@@ -90,6 +90,7 @@ class MigrateUniqueIDFlow(RepairsFlow):
                     config_entry,
                     unique_id=self.description_placeholders["new_unique_id"],
                 )
+                self.hass.config_entries.async_schedule_reload(config_entry.entry_id)
             return self.async_create_entry(data={})
 
         return self.async_show_form(

tests/components/zwave_js/test_config_flow.py

@@ -883,9 +883,9 @@ async def test_usb_discovery_migration(
     addon_options["device"] = "/dev/ttyUSB0"
     entry = integration
     assert client.connect.call_count == 1
+    assert entry.unique_id == "3245146787"
     hass.config_entries.async_update_entry(
         entry,
-        unique_id="1234",
         data={
             "url": "ws://localhost:3000",
             "use_addon": True,
@@ -893,6 +893,11 @@ async def test_usb_discovery_migration(
         },
     )
 
+    async def mock_restart_addon(addon_slug: str) -> None:
+        client.driver.controller.data["homeId"] = 1234
+
+    restart_addon.side_effect = mock_restart_addon
+
     async def mock_backup_nvm_raw():
         await asyncio.sleep(0)
         client.driver.controller.emit(
@@ -914,6 +919,7 @@ async def test_usb_discovery_migration(
             "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"}
         )
@@ -967,7 +973,7 @@ async def test_usb_discovery_migration(
     assert restart_addon.call_args == call("core_zwave_js")
 
     version_info = get_server_version.return_value
-    version_info.home_id = 5678
+    version_info.home_id = 3245146787
 
     result = await hass.config_entries.flow.async_configure(result["flow_id"])
 
@@ -991,7 +997,7 @@ async def test_usb_discovery_migration(
     assert entry.data["usb_path"] == USB_DISCOVERY_INFO.device
     assert entry.data["use_addon"] is True
     assert "keep_old_devices" not in entry.data
-    assert entry.unique_id == "5678"
+    assert entry.unique_id == "3245146787"
 
 
 @pytest.mark.usefixtures("supervisor", "addon_running")
@@ -1008,9 +1014,9 @@ async def test_usb_discovery_migration_restore_driver_ready_timeout(
     addon_options["device"] = "/dev/ttyUSB0"
     entry = integration
     assert client.connect.call_count == 1
+    assert entry.unique_id == "3245146787"
     hass.config_entries.async_update_entry(
         entry,
-        unique_id="1234",
         data={
             "url": "ws://localhost:3000",
             "use_addon": True,
@@ -1018,6 +1024,11 @@ async def test_usb_discovery_migration_restore_driver_ready_timeout(
         },
     )
 
+    async def mock_restart_addon(addon_slug: str) -> None:
+        client.driver.controller.data["homeId"] = 1234
+
+    restart_addon.side_effect = mock_restart_addon
+
     async def mock_backup_nvm_raw():
         await asyncio.sleep(0)
         client.driver.controller.emit(
@@ -1039,6 +1050,7 @@ async def test_usb_discovery_migration_restore_driver_ready_timeout(
             "nvm restore progress",
             {"event": "nvm restore progress", "bytesWritten": 100, "total": 200},
         )
+        client.driver.controller.data["homeId"] = 3245146787
 
     client.driver.controller.async_restore_nvm = AsyncMock(side_effect=mock_restore_nvm)
 
@@ -1113,7 +1125,8 @@ async def test_usb_discovery_migration_restore_driver_ready_timeout(
     assert entry.data["url"] == "ws://host1:3001"
     assert entry.data["usb_path"] == USB_DISCOVERY_INFO.device
     assert entry.data["use_addon"] is True
-    assert "keep_old_devices" not in entry.data
+    assert entry.unique_id == "1234"
+    assert "keep_old_devices" in entry.data
 
 
 @pytest.mark.usefixtures("supervisor", "addon_installed")
@@ -3011,6 +3024,7 @@ async def test_reconfigure_different_device(
     entry = integration
     data = {**entry.data, **entry_data}
     hass.config_entries.async_update_entry(entry, data=data, unique_id="1234")
+    client.driver.controller.data["homeId"] = 1234
 
     assert entry.data["url"] == "ws://test.org"
 
@@ -3164,6 +3178,7 @@ async def test_reconfigure_addon_restart_failed(
     entry = integration
     data = {**entry.data, **entry_data}
     hass.config_entries.async_update_entry(entry, data=data, unique_id="1234")
+    client.driver.controller.data["homeId"] = 1234
 
     assert entry.data["url"] == "ws://test.org"
 
@@ -3554,10 +3569,12 @@ async def test_reconfigure_migrate_low_sdk_version(
     (
         "restore_server_version_side_effect",
         "final_unique_id",
+        "keep_old_devices",
+        "device_entry_count",
     ),
     [
-        (None, "3245146787"),
+        (None, "3245146787", False, 2),
-        (aiohttp.ClientError("Boom"), "5678"),
+        (aiohttp.ClientError("Boom"), "5678", True, 4),
     ],
 )
 async def test_reconfigure_migrate_with_addon(
@@ -3572,12 +3589,15 @@ async def test_reconfigure_migrate_with_addon(
     get_server_version: AsyncMock,
     restore_server_version_side_effect: Exception | None,
     final_unique_id: str,
+    keep_old_devices: bool,
+    device_entry_count: int,
 ) -> None:
     """Test migration flow with add-on."""
     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={
@@ -3745,10 +3765,10 @@ async def test_reconfigure_migrate_with_addon(
     assert entry.data["url"] == "ws://host1:3001"
     assert entry.data["usb_path"] == "/test"
     assert entry.data["use_addon"] is True
-    assert "keep_old_devices" not in entry.data
+    assert ("keep_old_devices" in entry.data) is keep_old_devices
     assert entry.unique_id == final_unique_id
 
-    assert len(device_registry.devices) == 2
+    assert len(device_registry.devices) == device_entry_count
     controller_device_id_ext = (
         f"{controller_device_id}-{controller_node.manufacturer_id}:"
         f"{controller_node.product_type}:{controller_node.product_id}"
@@ -3780,9 +3800,10 @@ async def test_reconfigure_migrate_restore_driver_ready_timeout(
     """Test migration flow with driver ready timeout after nvm restore."""
     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,
-        unique_id="1234",
         data={
             "url": "ws://localhost:3000",
             "use_addon": True,
@@ -3790,6 +3811,11 @@ async def test_reconfigure_migrate_restore_driver_ready_timeout(
         },
     )
 
+    async def mock_restart_addon(addon_slug: str) -> None:
+        client.driver.controller.data["homeId"] = 1234
+
+    restart_addon.side_effect = mock_restart_addon
+
     async def mock_backup_nvm_raw():
         await asyncio.sleep(0)
         client.driver.controller.emit(
@@ -3811,6 +3837,7 @@ async def test_reconfigure_migrate_restore_driver_ready_timeout(
             "nvm restore progress",
             {"event": "nvm restore progress", "bytesWritten": 100, "total": 200},
         )
+        client.driver.controller.data["homeId"] = 3245146787
 
     client.driver.controller.async_restore_nvm = AsyncMock(side_effect=mock_restore_nvm)
 
@@ -3894,7 +3921,8 @@ async def test_reconfigure_migrate_restore_driver_ready_timeout(
     assert entry.data["url"] == "ws://host1:3001"
     assert entry.data["usb_path"] == "/test"
     assert entry.data["use_addon"] is True
-    assert "keep_old_devices" not in entry.data
+    assert "keep_old_devices" in entry.data
+    assert entry.unique_id == "1234"
 
 
 async def test_reconfigure_migrate_backup_failure(

tests/components/zwave_js/test_init.py

@@ -2070,12 +2070,8 @@ async def test_server_logging(
         await hass.config_entries.async_setup(entry.entry_id)
         await hass.async_block_till_done()
         assert len(client.async_send_command.call_args_list) == 2
-        assert client.async_send_command.call_args_list[0][0][0] == {
+        assert "driver.update_log_config" not in {
-            "command": "controller.get_provisioning_entries",
+            call[0][0]["command"] for call in client.async_send_command.call_args_list
-        }
-        assert client.async_send_command.call_args_list[1][0][0] == {
-            "command": "controller.get_provisioning_entry",
-            "dskOrNodeId": 1,
         }
         assert not client.enable_server_logging.called
         assert not client.disable_server_logging.called

tests/components/zwave_js/test_repairs.py

@@ -1,13 +1,14 @@
 """Test the Z-Wave JS repairs module."""
 
 from copy import deepcopy
-from unittest.mock import patch
+from unittest.mock import MagicMock, patch
 
 import pytest
 from zwave_js_server.event import Event
 from zwave_js_server.model.node import Node
 
 from homeassistant.components.zwave_js import DOMAIN
+from homeassistant.components.zwave_js.const import CONF_KEEP_OLD_DEVICES
 from homeassistant.components.zwave_js.helpers import get_device_id
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers import device_registry as dr, issue_registry as ir
@@ -276,8 +277,12 @@ async def test_migrate_unique_id(
     hass: HomeAssistant,
     hass_client: ClientSessionGenerator,
     hass_ws_client: WebSocketGenerator,
+    device_registry: dr.DeviceRegistry,
+    client: MagicMock,
+    multisensor_6: Node,
 ) -> None:
     """Test the migrate unique id flow."""
+    node = multisensor_6
     old_unique_id = "123456789"
     config_entry = MockConfigEntry(
         domain=DOMAIN,
@@ -289,8 +294,27 @@ async def test_migrate_unique_id(
     )
     config_entry.add_to_hass(hass)
 
+    # Remove the node from the current controller's known nodes.
+    client.driver.controller.nodes.pop(node.node_id)
+
+    # Create a device entry for the node connected to the old controller.
+    device_entry = device_registry.async_get_or_create(
+        config_entry_id=config_entry.entry_id,
+        identifiers={(DOMAIN, f"{old_unique_id}-{node.node_id}")},
+        name="Node connected to old controller",
+    )
+    assert device_entry.name == "Node connected to old controller"
+
     await hass.config_entries.async_setup(config_entry.entry_id)
 
+    assert CONF_KEEP_OLD_DEVICES in config_entry.data
+    assert config_entry.data[CONF_KEEP_OLD_DEVICES] is True
+    stored_devices = dr.async_entries_for_config_entry(
+        device_registry, config_entry.entry_id
+    )
+    assert len(stored_devices) == 2
+    assert device_entry.id in {device.id for device in stored_devices}
+
     await async_process_repairs_platforms(hass)
     ws_client = await hass_ws_client(hass)
     http_client = await hass_client()
@@ -317,6 +341,13 @@ async def test_migrate_unique_id(
 
     # Apply fix
     data = await process_repair_fix_flow(http_client, flow_id)
+    await hass.async_block_till_done()
+
+    stored_devices = dr.async_entries_for_config_entry(
+        device_registry, config_entry.entry_id
+    )
+    assert len(stored_devices) == 1
+    assert device_entry.id not in {device.id for device in stored_devices}
 
     assert data["type"] == "create_entry"
     assert config_entry.unique_id == "3245146787"