diff --git a/tests/components/esphome/test_entity.py b/tests/components/esphome/test_entity.py index d7edd6d096e..9fda478e79a 100644 --- a/tests/components/esphome/test_entity.py +++ b/tests/components/esphome/test_entity.py @@ -964,11 +964,13 @@ async def test_entity_switches_between_devices( ), ] - # 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) + # Update the entity info by changing what the mock returns + mock_client.list_entities_services = AsyncMock( + return_value=(updated_entity_info, []) + ) + # Trigger a reconnect to simulate the entity info update + await device.mock_disconnect(expected_disconnect=False) + await device.mock_connect() # Verify entity is now on sub device 1 sub_device_1 = device_registry.async_get_device( @@ -991,8 +993,11 @@ async def test_entity_switches_between_devices( ), ] - for callback_ in callbacks: - callback_(updated_entity_info) + mock_client.list_entities_services = AsyncMock( + return_value=(updated_entity_info, []) + ) + await device.mock_disconnect(expected_disconnect=False) + await device.mock_connect() # Verify entity is now on sub device 2 sub_device_2 = device_registry.async_get_device( @@ -1015,10 +1020,91 @@ async def test_entity_switches_between_devices( ), ] - for callback_ in callbacks: - callback_(updated_entity_info) + mock_client.list_entities_services = AsyncMock( + return_value=(updated_entity_info, []) + ) + await device.mock_disconnect(expected_disconnect=False) + await device.mock_connect() # 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 + + +async def test_entity_switches_devices_entity_not_in_registry( + hass: HomeAssistant, + entity_registry: er.EntityRegistry, + device_registry: dr.DeviceRegistry, + mock_client: APIClient, + mock_esphome_device: MockESPHomeDeviceType, +) -> None: + """Test entity device switch when entity is not in registry.""" + # 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 on sub device 1 + entity_info = [ + BinarySensorInfo( + object_id="sensor", + key=1, + name="Test Sensor", + unique_id="sensor", + device_id=11111111, + ), + ] + + 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 created 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 + + # Remove the entity from registry + entity_registry.async_remove(sensor_entity.entity_id) + + # Verify it's removed + assert entity_registry.async_get("binary_sensor.test_sensor") is None + + # Try to update entity to move to sub device 2 + updated_entity_info = [ + BinarySensorInfo( + object_id="sensor", + key=1, + name="Test Sensor", + unique_id="sensor", + device_id=22222222, # Try to move to sub device 2 + ), + ] + + # Update should not fail even though entity is not in registry + mock_client.list_entities_services = AsyncMock( + return_value=(updated_entity_info, []) + ) + await device.mock_disconnect(expected_disconnect=False) + await device.mock_connect() + + # Entity should remain not in registry + assert entity_registry.async_get("binary_sensor.test_sensor") is None