Don't remove zwave_js devices automatically

homeassistant/components/zwave_js/__init__.py

@@ -282,7 +282,8 @@ async def handle_logging_changed(_: Event | None = None) -> None:
             for node in controller.nodes.values()
         ]
 
-        # Devices that are in the device registry that are not known by the controller can be removed
+        # Devices that are in the device registry that are not known by the controller
+        # can be removed
         for device in stored_devices:
             if device not in known_devices:
                 self.dev_reg.async_remove_device(device.id)
@@ -509,25 +510,46 @@ def register_node_in_dev_reg(self, node: ZwaveNode) -> dr.DeviceEntry:
         driver = self.driver_events.driver
         device_id = get_device_id(driver, node)
         device_id_ext = get_device_id_ext(driver, node)
-        device = self.dev_reg.async_get_device(identifiers={device_id})
+        node_id_device = self.dev_reg.async_get_device(identifiers={device_id})
         via_device_id = None
         controller = driver.controller
         # Get the controller node device ID if this node is not the controller
         if controller.own_node and controller.own_node != node:
             via_device_id = get_device_id(driver, controller.own_node)
 
-        # Replace the device if it can be determined that this node is not the
-        # same product as it was previously.
-        if (
-            device_id_ext
-            and device
-            and len(device.identifiers) == 2
-            and device_id_ext not in device.identifiers
-        ):
-            self.remove_device(device)
-            device = None
-
         if device_id_ext:
+            # If there is a device with this node ID but with a different hardware
+            # signature, remove the node ID based identifier from it. The hardware
+            # signature can be different for one of two reasons: 1) in the ideal
+            # scenario, the node was replaced with a different node that's a different
+            # device entirely, or 2) the device erroneously advertised the wrong
+            # hardware identifiers (this is known to happen due to poor RF conditions).
+            # While we would like to remove the old device automatically for case 1, we
+            # have no way to distinguish between these reasons so we leave it up to the
+            # user to remove the old device manually.
+            if (
+                node_id_device
+                and len(node_id_device.identifiers) == 2
+                and device_id_ext not in node_id_device.identifiers
+            ):
+                new_identifiers = node_id_device.identifiers.copy()
+                new_identifiers.remove(device_id)
+                self.dev_reg.async_update_device(
+                    node_id_device.id, new_identifiers=new_identifiers
+                )
+            # If there is an orphaned device that already exists with this hardware
+            # based identifier, add the node ID based identifier to the orphaned
+            # device.
+            if (
+                hardware_device := self.dev_reg.async_get_device(
+                    identifiers={device_id_ext}
+                )
+            ) and len(hardware_device.identifiers) == 1:
+                new_identifiers = hardware_device.identifiers.copy()
+                new_identifiers.add(device_id)
+                self.dev_reg.async_update_device(
+                    hardware_device.id, new_identifiers=new_identifiers
+                )
             ids = {device_id, device_id_ext}
         else:
             ids = {device_id}
@@ -769,9 +791,12 @@ def async_on_notification(self, event: dict[str, Any]) -> None:
             return
 
         driver = self.controller_events.driver_events.driver
-        notification: EntryControlNotification | NotificationNotification | PowerLevelNotification | MultilevelSwitchNotification = event[
-            "notification"
-        ]
+        notification: (
+            EntryControlNotification
+            | NotificationNotification
+            | PowerLevelNotification
+            | MultilevelSwitchNotification
+        ) = event["notification"]
         device = self.dev_reg.async_get_device(
             identifiers={get_device_id(driver, notification.node)}
         )
@@ -984,6 +1009,39 @@ async def async_remove_entry(hass: HomeAssistant, entry: ConfigEntry) -> None:
         LOGGER.error(err)
 
 
+async def async_remove_config_entry_device(
+    hass: HomeAssistant, config_entry: ConfigEntry, device_entry: dr.DeviceEntry
+) -> bool:
+    """Remove a config entry from a device."""
+    entry_hass_data = hass.data[DOMAIN][config_entry.entry_id]
+    client: ZwaveClient = entry_hass_data[DATA_CLIENT]
+
+    # Driver may not be ready yet so we can't allow users to remove a device since
+    # we need to check if the device is still known to the controller
+    if (driver := client.driver) is None:
+        LOGGER.error("Driver for %s is not ready", config_entry.title)
+        return False
+
+    # If a node is found on the controller that matches the hardware based identifier
+    # on the device, prevent the device from being removed.
+    if next(
+        (
+            node
+            for node in driver.controller.nodes.values()
+            if get_device_id_ext(driver, node) in device_entry.identifiers
+        ),
+        None,
+    ):
+        return False
+
+    controller_events: ControllerEvents = entry_hass_data[
+        DATA_DRIVER_EVENTS
+    ].controller_events
+    controller_events.registered_unique_ids.pop(device_entry.id, None)
+    controller_events.discovered_value_ids.pop(device_entry.id, None)
+    return True
+
+
 async def async_ensure_addon_running(hass: HomeAssistant, entry: ConfigEntry) -> None:
     """Ensure that Z-Wave JS add-on is installed and running."""
     addon_manager = _get_addon_manager(hass)

tests/components/zwave_js/test_init.py

@@ -30,6 +30,7 @@
 from .common import AIR_TEMPERATURE_SENSOR, EATON_RF9640_ENTITY
 
 from tests.common import MockConfigEntry, async_get_persistent_notifications
+from tests.typing import WebSocketGenerator
 
 
 @pytest.fixture(name="connect_timeout")
@@ -962,6 +963,7 @@ async def test_removed_device(
     # Make sure there are the same number of devices
     dev_reg = dr.async_get(hass)
     device_entries = dr.async_entries_for_config_entry(dev_reg, integration.entry_id)
+    logging.getLogger(__name__).error(device_entries)
     assert len(device_entries) == 3
 
     # Check how many entities there are
@@ -1016,6 +1018,7 @@ async def test_node_removed(
     client.driver.controller.receive_event(Event("node added", event))
     await hass.async_block_till_done()
     old_device = dev_reg.async_get_device(identifiers={(DOMAIN, device_id)})
+    assert old_device
     assert old_device.id
 
     event = {"node": node, "reason": 0}
@@ -1139,6 +1142,7 @@ async def test_replace_different_node(
     hank_binary_switch_state,
     client,
     integration,
+    hass_ws_client: WebSocketGenerator,
 ) -> None:
     """Test when a node is replaced with a different node."""
     dev_reg = dr.async_get(hass)
@@ -1147,11 +1151,11 @@ async def test_replace_different_node(
     state["nodeId"] = node_id
 
     device_id = f"{client.driver.controller.home_id}-{node_id}"
-    multisensor_6_device_id = (
+    multisensor_6_device_id_ext = (
         f"{device_id}-{multisensor_6.manufacturer_id}:"
         f"{multisensor_6.product_type}:{multisensor_6.product_id}"
     )
-    hank_device_id = (
+    hank_device_id_ext = (
         f"{device_id}-{state['manufacturerId']}:"
         f"{state['productType']}:"
         f"{state['productId']}"
@@ -1160,16 +1164,15 @@ async def test_replace_different_node(
     device = dev_reg.async_get_device(identifiers={(DOMAIN, device_id)})
     assert device
     assert device == dev_reg.async_get_device(
-        identifiers={(DOMAIN, multisensor_6_device_id)}
+        identifiers={(DOMAIN, multisensor_6_device_id_ext)}
     )
     assert device.manufacturer == "AEON Labs"
     assert device.model == "ZW100"
     dev_id = device.id
 
     assert hass.states.get(AIR_TEMPERATURE_SENSOR)
 
-    # A replace node event has the extra field "replaced" set to True
-    # to distinguish it from an exclusion
+    # Remove existing node
     event = Event(
         type="node removed",
         data={
@@ -1183,8 +1186,11 @@ async def test_replace_different_node(
     await hass.async_block_till_done()
 
     # Device should still be there after the node was removed
-    device = dev_reg.async_get(dev_id)
+    device = dev_reg.async_get_device(
+        identifiers={(DOMAIN, multisensor_6_device_id_ext)}
+    )
     assert device
+    assert len(device.identifiers) == 2
 
     # When the node is replaced, a non-ready node added event is emitted
     event = Event(
@@ -1238,18 +1244,168 @@ async def test_replace_different_node(
     client.driver.receive_event(event)
     await hass.async_block_till_done()
 
-    # Old device and entities were removed, but the ID is re-used
-    device = dev_reg.async_get(dev_id)
-    assert device
-    assert device == dev_reg.async_get_device(identifiers={(DOMAIN, device_id)})
-    assert device == dev_reg.async_get_device(identifiers={(DOMAIN, hank_device_id)})
-    assert not dev_reg.async_get_device(identifiers={(DOMAIN, multisensor_6_device_id)})
-    assert device.manufacturer == "HANK Electronics Ltd."
-    assert device.model == "HKZW-SO01"
+    # node ID based device identifier should be moved from the old multisensor device
+    # to the new hank device and both the old and new devices should exist.
+    new_device = dev_reg.async_get_device(identifiers={(DOMAIN, device_id)})
+    assert new_device
+    hank_device = dev_reg.async_get_device(identifiers={(DOMAIN, hank_device_id_ext)})
+    assert hank_device
+    assert hank_device == new_device
+    assert hank_device.identifiers == {
+        (DOMAIN, device_id),
+        (DOMAIN, hank_device_id_ext),
+    }
+    multisensor_6_device = dev_reg.async_get_device(
+        identifiers={(DOMAIN, multisensor_6_device_id_ext)}
+    )
+    assert multisensor_6_device
+    assert multisensor_6_device != new_device
+    assert multisensor_6_device.identifiers == {(DOMAIN, multisensor_6_device_id_ext)}
+
+    assert new_device.manufacturer == "HANK Electronics Ltd."
+    assert new_device.model == "HKZW-SO01"
 
-    assert not hass.states.get(AIR_TEMPERATURE_SENSOR)
+    # We keep the old entities in case there are customizations that a user wants to
+    # keep. They can always delete the device and that will remove the entities as well.
+    assert hass.states.get(AIR_TEMPERATURE_SENSOR)
     assert hass.states.get("switch.smart_plug_with_two_usb_ports")
 
+    # Try to add back the first node to see if the device IDs are correct
+
+    # Remove existing node
+    event = Event(
+        type="node removed",
+        data={
+            "source": "controller",
+            "event": "node removed",
+            "reason": 3,
+            "node": state,
+        },
+    )
+    client.driver.receive_event(event)
+    await hass.async_block_till_done()
+
+    # Device should still be there after the node was removed
+    device = dev_reg.async_get_device(identifiers={(DOMAIN, hank_device_id_ext)})
+    assert device
+    assert len(device.identifiers) == 2
+
+    # When the node is replaced, a non-ready node added event is emitted
+    event = Event(
+        type="node added",
+        data={
+            "source": "controller",
+            "event": "node added",
+            "node": {
+                "nodeId": multisensor_6.node_id,
+                "index": 0,
+                "status": 4,
+                "ready": False,
+                "isSecure": False,
+                "interviewAttempts": 1,
+                "endpoints": [
+                    {"nodeId": multisensor_6.node_id, "index": 0, "deviceClass": None}
+                ],
+                "values": [],
+                "deviceClass": None,
+                "commandClasses": [],
+                "interviewStage": "None",
+                "statistics": {
+                    "commandsTX": 0,
+                    "commandsRX": 0,
+                    "commandsDroppedRX": 0,
+                    "commandsDroppedTX": 0,
+                    "timeoutResponse": 0,
+                },
+                "isControllerNode": False,
+            },
+            "result": {},
+        },
+    )
+
+    client.driver.receive_event(event)
+    await hass.async_block_till_done()
+
+    # Mark node as ready
+    event = Event(
+        type="ready",
+        data={
+            "source": "node",
+            "event": "ready",
+            "nodeId": node_id,
+            "nodeState": multisensor_6_state,
+        },
+    )
+    client.driver.receive_event(event)
+    await hass.async_block_till_done()
+
+    assert await async_setup_component(hass, "config", {})
+
+    # node ID based device identifier should be moved from the new hank device
+    # to the old multisensor device and both the old and new devices should exist.
+    old_device = dev_reg.async_get_device(identifiers={(DOMAIN, device_id)})
+    assert old_device
+    hank_device = dev_reg.async_get_device(identifiers={(DOMAIN, hank_device_id_ext)})
+    assert hank_device
+    assert hank_device != old_device
+    assert hank_device.identifiers == {(DOMAIN, hank_device_id_ext)}
+    multisensor_6_device = dev_reg.async_get_device(
+        identifiers={(DOMAIN, multisensor_6_device_id_ext)}
+    )
+    assert multisensor_6_device
+    assert multisensor_6_device == old_device
+    assert multisensor_6_device.identifiers == {
+        (DOMAIN, device_id),
+        (DOMAIN, multisensor_6_device_id_ext),
+    }
+
+    ws_client = await hass_ws_client(hass)
+
+    # Simulate the driver not being ready to ensure that the device removal handler
+    # does not crash
+    driver = client.driver
+    client.driver = None
+
+    await ws_client.send_json(
+        {
+            "id": 1,
+            "type": "config/device_registry/remove_config_entry",
+            "config_entry_id": integration.entry_id,
+            "device_id": hank_device.id,
+        }
+    )
+    response = await ws_client.receive_json()
+    assert not response["success"]
+
+    client.driver = driver
+
+    # Attempting to remove the hank device should pass, but removing the multisensor should not
+    await ws_client.send_json(
+        {
+            "id": 2,
+            "type": "config/device_registry/remove_config_entry",
+            "config_entry_id": integration.entry_id,
+            "device_id": hank_device.id,
+        }
+    )
+    response = await ws_client.receive_json()
+    assert response["success"]
+
+    await ws_client.send_json(
+        {
+            "id": 3,
+            "type": "config/device_registry/remove_config_entry",
+            "config_entry_id": integration.entry_id,
+            "device_id": multisensor_6_device.id,
+        }
+    )
+    response = await ws_client.receive_json()
+    assert not response["success"]
+    # assert await async_remove_config_entry_device(hass, integration, hank_device)
+    # assert not await async_remove_config_entry_device(
+    #     hass, integration, multisensor_6_device
+    # )
+
 
 async def test_node_model_change(
     hass: HomeAssistant, zp3111, client, integration