core/tests/helpers/test_device.py

"""Tests for the Device Utils."""

from unittest.mock import patch

import pytest
import voluptuous as vol

from homeassistant.core import HomeAssistant
from homeassistant.helpers import device_registry as dr, entity_registry as er
from homeassistant.helpers.device import (
    async_device_info_to_link_from_device_id,
    async_device_info_to_link_from_entity,
    async_entity_id_to_device,
    async_entity_id_to_device_id,
    async_remove_stale_devices_links_keep_current_device,
    async_remove_stale_devices_links_keep_entity_device,
)

from tests.common import MockConfigEntry


async def test_entity_id_to_device_device_id(
    hass: HomeAssistant,
    device_registry: dr.DeviceRegistry,
    entity_registry: er.EntityRegistry,
) -> None:
    """Test returning an entity's device / device ID."""
    config_entry = MockConfigEntry(domain="my")
    config_entry.add_to_hass(hass)

    device = device_registry.async_get_or_create(
        identifiers={("test", "current_device")},
        connections={("mac", "30:31:32:33:34:00")},
        config_entry_id=config_entry.entry_id,
    )
    assert device is not None

    # Entity registry
    entity = entity_registry.async_get_or_create(
        "sensor",
        "test",
        "source",
        config_entry=config_entry,
        device_id=device.id,
    )
    await hass.async_block_till_done()
    assert entity_registry.async_get("sensor.test_source") is not None

    device_id = async_entity_id_to_device_id(
        hass,
        entity_id_or_uuid=entity.entity_id,
    )
    assert device_id == device.id
    assert (
        async_entity_id_to_device(
            hass,
            entity_id_or_uuid=entity.entity_id,
        )
        == device
    )

    assert (
        async_entity_id_to_device_id(
            hass,
            entity_id_or_uuid="unknown.entity_id",
        )
        is None
    )
    assert (
        async_entity_id_to_device(
            hass,
            entity_id_or_uuid="unknown.entity_id",
        )
        is None
    )

    device_id = async_entity_id_to_device_id(
        hass,
        entity_id_or_uuid=entity.id,
    )
    assert device_id == device.id
    assert (
        async_entity_id_to_device(
            hass,
            entity_id_or_uuid=entity.id,
        )
        == device
    )

    with pytest.raises(vol.Invalid):
        async_entity_id_to_device_id(
            hass,
            entity_id_or_uuid="unknown_uuid",
        )

    with pytest.raises(vol.Invalid):
        async_entity_id_to_device(
            hass,
            entity_id_or_uuid="unknown_uuid",
        )


async def test_device_info_to_link(
    hass: HomeAssistant,
    device_registry: dr.DeviceRegistry,
    entity_registry: er.EntityRegistry,
) -> None:
    """The link helpers are deprecated and always return None.

    A device_info carrying another device's identifiers implicitly added the caller's
    config entry to that device, which a single-config-entry device can't represent - it
    would silently fork a duplicate instead. Entities still attach to another config
    entry's device by setting entity.device_entry.
    """
    config_entry = MockConfigEntry(domain="my")
    config_entry.add_to_hass(hass)

    device = device_registry.async_get_or_create(
        identifiers={("test", "my_device")},
        connections={("mac", "30:31:32:33:34:00")},
        config_entry_id=config_entry.entry_id,
    )

    # Source entity registry
    source_entity = entity_registry.async_get_or_create(
        "sensor",
        "test",
        "source",
        config_entry=config_entry,
        device_id=device.id,
    )
    await hass.async_block_till_done()
    assert entity_registry.async_get("sensor.test_source") is not None

    # No link device_info is returned, even for an existing entity and device
    with patch("homeassistant.helpers.device.report_usage") as report_usage:
        assert (
            async_device_info_to_link_from_entity(
                hass, entity_id_or_uuid=source_entity.entity_id
            )
            is None
        )
        assert (
            async_device_info_to_link_from_device_id(hass, device_id=device.id) is None
        )
    assert report_usage.call_count == 2

    # With a non-existent entity id
    assert (
        async_device_info_to_link_from_entity(hass, entity_id_or_uuid="sensor.invalid")
        is None
    )

    # With a non-existent device id
    assert async_device_info_to_link_from_device_id(hass, device_id="abcdefghi") is None

    # With a None device id
    assert async_device_info_to_link_from_device_id(hass, device_id=None) is None


async def test_remove_stale_device_links_keep_entity_device(
    hass: HomeAssistant,
    device_registry: dr.DeviceRegistry,
    entity_registry: er.EntityRegistry,
) -> None:
    """Test cleaning works for entity."""
    helper_config_entry = MockConfigEntry(domain="helper_integration")
    helper_config_entry.add_to_hass(hass)
    host_config_entry = MockConfigEntry(domain="host_integration")
    host_config_entry.add_to_hass(hass)

    current_device = device_registry.async_get_or_create(
        identifiers={("test", "current_device")},
        connections={("mac", "30:31:32:33:34:00")},
        config_entry_id=helper_config_entry.entry_id,
    )

    stale_device_1 = device_registry.async_get_or_create(
        identifiers={("test", "stale_device_1")},
        connections={("mac", "30:31:32:33:34:01")},
        config_entry_id=helper_config_entry.entry_id,
    )

    device_registry.async_get_or_create(
        identifiers={("test", "stale_device_2")},
        connections={("mac", "30:31:32:33:34:02")},
        config_entry_id=helper_config_entry.entry_id,
    )

    # Source entity
    source_entity = entity_registry.async_get_or_create(
        "sensor",
        "host_integration",
        "source",
        config_entry=host_config_entry,
        device_id=current_device.id,
    )
    assert entity_registry.async_get(source_entity.entity_id) is not None

    # Helper entity connected to a stale device
    helper_entity = entity_registry.async_get_or_create(
        "sensor",
        "helper_integration",
        "helper",
        config_entry=helper_config_entry,
        device_id=stale_device_1.id,
    )
    assert entity_registry.async_get(helper_entity.entity_id) is not None

    devices_helper_entry = device_registry.devices.get_devices_for_config_entry_id(
        helper_config_entry.entry_id
    )

    # 3 devices linked to the config entry are expected (1 current device + 2 stales)
    assert len(devices_helper_entry) == 3

    # Manual cleanup should unlink stale devices from the config entry
    async_remove_stale_devices_links_keep_entity_device(
        hass,
        entry_id=helper_config_entry.entry_id,
        source_entity_id_or_uuid=source_entity.entity_id,
    )

    await hass.async_block_till_done()

    devices_helper_entry = device_registry.devices.get_devices_for_config_entry_id(
        helper_config_entry.entry_id
    )

    # After cleanup, only one device is expected to be linked to the config entry, and
    # the entities should exist and be linked to the current device
    assert len(devices_helper_entry) == 1
    assert current_device in devices_helper_entry
    assert entity_registry.async_get(source_entity.entity_id) is not None
    assert entity_registry.async_get(helper_entity.entity_id) is not None


async def test_remove_stale_devices_links_keep_current_device(
    hass: HomeAssistant,
    device_registry: dr.DeviceRegistry,
) -> None:
    """Test cleanup works for device id."""
    config_entry = MockConfigEntry(domain="hue")
    config_entry.add_to_hass(hass)

    current_device = device_registry.async_get_or_create(
        identifiers={("test", "current_device")},
        connections={("mac", "30:31:32:33:34:00")},
        config_entry_id=config_entry.entry_id,
    )
    assert current_device is not None

    device_registry.async_get_or_create(
        identifiers={("test", "stale_device_1")},
        connections={("mac", "30:31:32:33:34:01")},
        config_entry_id=config_entry.entry_id,
    )

    device_registry.async_get_or_create(
        identifiers={("test", "stale_device_2")},
        connections={("mac", "30:31:32:33:34:02")},
        config_entry_id=config_entry.entry_id,
    )

    devices_config_entry = device_registry.devices.get_devices_for_config_entry_id(
        config_entry.entry_id
    )

    # 3 devices linked to the config entry are expected (1 current device + 2 stales)
    assert len(devices_config_entry) == 3

    # Manual cleanup should unlink stales devices from the config entry
    async_remove_stale_devices_links_keep_current_device(
        hass,
        entry_id=config_entry.entry_id,
        current_device_id=current_device.id,
    )

    devices_config_entry = device_registry.devices.get_devices_for_config_entry_id(
        config_entry.entry_id
    )

    # After cleanup, only one device is expected to be linked to the config entry
    assert len(devices_config_entry) == 1

    assert current_device in devices_config_entry