handle device switching

homeassistant/components/esphome/entity.py

@@ -24,6 +24,7 @@ from homeassistant.helpers import (
     config_validation as cv,
     device_registry as dr,
     entity_platform,
+    entity_registry as er,
 )
 from homeassistant.helpers.device_registry import DeviceInfo
 from homeassistant.helpers.entity import Entity
@@ -53,21 +54,51 @@ def async_static_info_updated(
 ) -> None:
     """Update entities of this platform when entities are listed."""
     current_infos = entry_data.info[info_type]
+    device_info = entry_data.device_info
+    if TYPE_CHECKING:
+        assert device_info is not None
     new_infos: dict[int, EntityInfo] = {}
     add_entities: list[_EntityT] = []
 
+    ent_reg = er.async_get(hass)
+    dev_reg = dr.async_get(hass)
+
     for info in infos:
-        if not current_infos.pop(info.key, None):
-            # Create new entity
+        new_infos[info.key] = info
+
+        # Create new entity if it doesn't exist
+        if not (old_info := current_infos.pop(info.key, None)):
             entity = entity_type(entry_data, platform.domain, info, state_type)
             add_entities.append(entity)
-        new_infos[info.key] = info
+            continue
+
+        # Entity exists - check if device_id has changed
+        if old_info.device_id == info.device_id:
+            continue
+
+        # Entity has switched devices, update its device assignment
+        unique_id = build_unique_id(device_info.mac_address, info)
+        entity_id = ent_reg.async_get_entity_id(platform.domain, DOMAIN, unique_id)
+        if not entity_id:
+            continue
+
+        # Determine the new device
+        if info.device_id:
+            # Entity now belongs to a sub device
+            new_device = dev_reg.async_get_device(
+                identifiers={(DOMAIN, f"{device_info.mac_address}_{info.device_id}")}
+            )
+        else:
+            # Entity now belongs to the main device
+            new_device = dev_reg.async_get_device(
+                connections={(dr.CONNECTION_NETWORK_MAC, device_info.mac_address)}
+            )
+
+        if new_device:
+            ent_reg.async_update_entity(entity_id, device_id=new_device.id)
 
     # Anything still in current_infos is now gone
     if current_infos:
-        device_info = entry_data.device_info
-        if TYPE_CHECKING:
-            assert device_info is not None
         entry_data.async_remove_entities(
             hass, current_infos.values(), device_info.mac_address
         )

tests/components/esphome/test_entity.py

@@ -901,3 +901,124 @@ async def test_entity_friendly_names_with_empty_device_names(
     state_4 = hass.states.get("binary_sensor.test_main_status")
     assert state_4 is not None
     assert state_4.attributes[ATTR_FRIENDLY_NAME] == "Main Device Main Status"
+
+
+async def test_entity_switches_between_devices(
+    hass: HomeAssistant,
+    entity_registry: er.EntityRegistry,
+    device_registry: dr.DeviceRegistry,
+    mock_client: APIClient,
+    mock_esphome_device: MockESPHomeDeviceType,
+) -> None:
+    """Test that entities can switch between devices correctly."""
+    # Define sub devices
+    sub_devices = [
+        SubDeviceInfo(device_id=11111111, name="Sub Device 1", area_id=0),
+        SubDeviceInfo(device_id=22222222, name="Sub Device 2", area_id=0),
+    ]
+
+    device_info = {
+        "devices": sub_devices,
+    }
+
+    # Create initial entity assigned to main device (no device_id)
+    entity_info = [
+        BinarySensorInfo(
+            object_id="sensor",
+            key=1,
+            name="Test Sensor",
+            unique_id="sensor",
+            # device_id omitted - entity belongs to main device
+        ),
+    ]
+
+    states = [
+        BinarySensorState(key=1, state=True, missing_state=False),
+    ]
+
+    device = await mock_esphome_device(
+        mock_client=mock_client,
+        device_info=device_info,
+        entity_info=entity_info,
+        states=states,
+    )
+
+    # Verify entity is on main device
+    main_device = device_registry.async_get_device(
+        connections={(dr.CONNECTION_NETWORK_MAC, device.device_info.mac_address)}
+    )
+    assert main_device is not None
+
+    sensor_entity = entity_registry.async_get("binary_sensor.test_sensor")
+    assert sensor_entity is not None
+    assert sensor_entity.device_id == main_device.id
+
+    # Test 1: Main device → Sub device 1
+    updated_entity_info = [
+        BinarySensorInfo(
+            object_id="sensor",
+            key=1,
+            name="Test Sensor",
+            unique_id="sensor",
+            device_id=11111111,  # Now on sub device 1
+        ),
+    ]
+
+    # Update the entity info (this would normally come from the ESP device)
+    entry_data = device.entry.runtime_data
+    callbacks = entry_data.entity_info_callbacks.get(BinarySensorInfo, [])
+    for callback_ in callbacks:
+        callback_(updated_entity_info)
+
+    # Verify entity is now on sub device 1
+    sub_device_1 = device_registry.async_get_device(
+        identifiers={(DOMAIN, f"{device.device_info.mac_address}_11111111")}
+    )
+    assert sub_device_1 is not None
+
+    sensor_entity = entity_registry.async_get("binary_sensor.test_sensor")
+    assert sensor_entity is not None
+    assert sensor_entity.device_id == sub_device_1.id
+
+    # Test 2: Sub device 1 → Sub device 2
+    updated_entity_info = [
+        BinarySensorInfo(
+            object_id="sensor",
+            key=1,
+            name="Test Sensor",
+            unique_id="sensor",
+            device_id=22222222,  # Now on sub device 2
+        ),
+    ]
+
+    for callback_ in callbacks:
+        callback_(updated_entity_info)
+
+    # Verify entity is now on sub device 2
+    sub_device_2 = device_registry.async_get_device(
+        identifiers={(DOMAIN, f"{device.device_info.mac_address}_22222222")}
+    )
+    assert sub_device_2 is not None
+
+    sensor_entity = entity_registry.async_get("binary_sensor.test_sensor")
+    assert sensor_entity is not None
+    assert sensor_entity.device_id == sub_device_2.id
+
+    # Test 3: Sub device 2 → Main device
+    updated_entity_info = [
+        BinarySensorInfo(
+            object_id="sensor",
+            key=1,
+            name="Test Sensor",
+            unique_id="sensor",
+            # device_id omitted - back to main device
+        ),
+    ]
+
+    for callback_ in callbacks:
+        callback_(updated_entity_info)
+
+    # Verify entity is back on main device
+    sensor_entity = entity_registry.async_get("binary_sensor.test_sensor")
+    assert sensor_entity is not None
+    assert sensor_entity.device_id == main_device.id