core/tests/components/waqi/test_init.py

"""Test the World Air Quality Index (WAQI) initialization."""

from typing import Any
from unittest.mock import AsyncMock, patch

from aiowaqi import WAQIError
import attr
import pytest

from homeassistant.components.waqi import DOMAIN
from homeassistant.components.waqi.const import CONF_STATION_NUMBER
from homeassistant.config_entries import ConfigEntryDisabler, ConfigEntryState
from homeassistant.const import CONF_API_KEY
from homeassistant.core import HomeAssistant
from homeassistant.helpers import device_registry as dr, entity_registry as er
from homeassistant.helpers.device_registry import DeviceEntryDisabler
from homeassistant.helpers.entity_registry import RegistryEntryDisabler

from . import setup_integration

from tests.common import MockConfigEntry


async def test_setup_failed(
    hass: HomeAssistant,
    mock_config_entry: MockConfigEntry,
    mock_waqi: AsyncMock,
) -> None:
    """Test setup failure due to API error."""
    mock_waqi.get_by_station_number.side_effect = WAQIError("API error")
    await setup_integration(hass, mock_config_entry)

    assert mock_config_entry.state is ConfigEntryState.SETUP_RETRY


@pytest.mark.usefixtures("mock_setup_entry")
async def test_migration_from_v1(
    hass: HomeAssistant,
    device_registry: dr.DeviceRegistry,
    entity_registry: er.EntityRegistry,
) -> None:
    """Test migration from version 1 to version 2."""
    # Create a v1 config entry with conversation options and an entity
    mock_config_entry = MockConfigEntry(
        domain=DOMAIN,
        data={CONF_API_KEY: "1234", CONF_STATION_NUMBER: 4584},
        version=1,
        unique_id="4584",
        title="de Jongweg, Utrecht",
    )
    mock_config_entry.add_to_hass(hass)
    mock_config_entry_2 = MockConfigEntry(
        domain=DOMAIN,
        data={CONF_API_KEY: "1234", CONF_STATION_NUMBER: 4585},
        version=1,
        unique_id="4585",
        title="Not de Jongweg, Utrecht",
    )
    mock_config_entry_2.add_to_hass(hass)

    device_1 = device_registry.async_get_or_create(
        config_entry_id=mock_config_entry.entry_id,
        identifiers={(DOMAIN, "4584")},
        name="de Jongweg, Utrecht",
        entry_type=dr.DeviceEntryType.SERVICE,
    )
    entity_registry.async_get_or_create(
        "sensor",
        DOMAIN,
        "4584_air_quality",
        config_entry=mock_config_entry,
        device_id=device_1.id,
        suggested_object_id="de_jongweg_utrecht",
    )

    device_2 = device_registry.async_get_or_create(
        config_entry_id=mock_config_entry_2.entry_id,
        identifiers={(DOMAIN, "4585")},
        name="Not de Jongweg, Utrecht",
        entry_type=dr.DeviceEntryType.SERVICE,
    )
    entity_registry.async_get_or_create(
        "sensor",
        DOMAIN,
        "4585_air_quality",
        config_entry=mock_config_entry_2,
        device_id=device_2.id,
        suggested_object_id="not_de_jongweg_utrecht",
    )

    await hass.config_entries.async_setup(mock_config_entry.entry_id)
    await hass.async_block_till_done()

    entries = hass.config_entries.async_entries(DOMAIN)
    assert len(entries) == 1
    entry = entries[0]
    assert entry.version == 2
    assert entry.minor_version == 1
    assert not entry.options
    assert entry.title == "WAQI"
    assert len(entry.subentries) == 2

    subentry = list(entry.subentries.values())[0]
    assert subentry.subentry_type == "station"
    assert subentry.data[CONF_STATION_NUMBER] == 4584
    assert subentry.unique_id == "4584"
    assert subentry.title == "de Jongweg, Utrecht"

    entity = entity_registry.async_get("sensor.de_jongweg_utrecht")
    assert entity.unique_id == "4584_air_quality"
    assert entity.config_subentry_id == subentry.subentry_id
    assert entity.config_entry_id == entry.entry_id

    assert (device := device_registry.async_get_device(identifiers={(DOMAIN, "4584")}))
    assert device.identifiers == {(DOMAIN, "4584")}
    assert device.id == device_1.id
    assert device.config_entries == {mock_config_entry.entry_id}
    assert device.config_entries_subentries == {
        mock_config_entry.entry_id: {subentry.subentry_id}
    }

    subentry = list(entry.subentries.values())[1]
    assert subentry.subentry_type == "station"
    assert subentry.data[CONF_STATION_NUMBER] == 4585
    assert subentry.unique_id == "4585"
    assert subentry.title == "Not de Jongweg, Utrecht"

    entity = entity_registry.async_get("sensor.not_de_jongweg_utrecht")
    assert entity.unique_id == "4585_air_quality"
    assert entity.config_subentry_id == subentry.subentry_id
    assert entity.config_entry_id == entry.entry_id
    assert (device := device_registry.async_get_device(identifiers={(DOMAIN, "4585")}))
    assert device.identifiers == {(DOMAIN, "4585")}
    assert device.id == device_2.id
    assert device.config_entries == {mock_config_entry.entry_id}
    assert device.config_entries_subentries == {
        mock_config_entry.entry_id: {subentry.subentry_id}
    }


@pytest.mark.parametrize(
    (
        "config_entry_disabled_by",
        "merged_config_entry_disabled_by",
        "sensor_subentry_data",
        "main_config_entry",
    ),
    [
        (
            [ConfigEntryDisabler.USER, None],
            None,
            [
                {
                    "sensor_entity_id": (
                        "sensor.not_de_jongweg_utrecht_air_quality_index"
                    ),
                    "device_disabled_by": None,
                    "entity_disabled_by": None,
                    "device": 1,
                },
                {
                    "sensor_entity_id": ("sensor.de_jongweg_utrecht_air_quality_index"),
                    "device_disabled_by": DeviceEntryDisabler.USER,
                    "entity_disabled_by": RegistryEntryDisabler.DEVICE,
                    "device": 0,
                },
            ],
            1,
        ),
        (
            [None, ConfigEntryDisabler.USER],
            None,
            [
                {
                    "sensor_entity_id": ("sensor.de_jongweg_utrecht_air_quality_index"),
                    "device_disabled_by": DeviceEntryDisabler.USER,
                    "entity_disabled_by": RegistryEntryDisabler.DEVICE,
                    "device": 0,
                },
                {
                    "sensor_entity_id": (
                        "sensor.not_de_jongweg_utrecht_air_quality_index"
                    ),
                    "device_disabled_by": None,
                    "entity_disabled_by": None,
                    "device": 1,
                },
            ],
            0,
        ),
        (
            [ConfigEntryDisabler.USER, ConfigEntryDisabler.USER],
            ConfigEntryDisabler.USER,
            [
                {
                    "sensor_entity_id": ("sensor.de_jongweg_utrecht_air_quality_index"),
                    "device_disabled_by": DeviceEntryDisabler.CONFIG_ENTRY,
                    "entity_disabled_by": RegistryEntryDisabler.CONFIG_ENTRY,
                    "device": 0,
                },
                {
                    "sensor_entity_id": (
                        "sensor.not_de_jongweg_utrecht_air_quality_index"
                    ),
                    # Device 2 was created enabled; the migration moves it onto the
                    # disabled merged config entry, so the move re-evaluates it as disabled
                    # by CONFIG_ENTRY (the entity keeps its own disabled_by - propagating a
                    # move-disable to entities is a separate mechanism)
                    "device_disabled_by": DeviceEntryDisabler.CONFIG_ENTRY,
                    "entity_disabled_by": None,
                    "device": 1,
                },
            ],
            0,
        ),
    ],
)
async def test_migration_from_v1_disabled(
    hass: HomeAssistant,
    device_registry: dr.DeviceRegistry,
    entity_registry: er.EntityRegistry,
    config_entry_disabled_by: list[ConfigEntryDisabler | None],
    merged_config_entry_disabled_by: ConfigEntryDisabler | None,
    sensor_subentry_data: list[dict[str, Any]],
    main_config_entry: int,
) -> None:
    """Test migration where the config entries are disabled."""
    mock_config_entry = MockConfigEntry(
        domain=DOMAIN,
        data={CONF_API_KEY: "1234", CONF_STATION_NUMBER: 4584},
        version=1,
        unique_id="4584",
        title="de Jongweg, Utrecht",
        disabled_by=config_entry_disabled_by[0],
    )
    mock_config_entry.add_to_hass(hass)
    mock_config_entry_2 = MockConfigEntry(
        domain=DOMAIN,
        data={CONF_API_KEY: "1234", CONF_STATION_NUMBER: 4585},
        version=1,
        unique_id="4585",
        title="Not de Jongweg, Utrecht",
        disabled_by=config_entry_disabled_by[1],
    )
    mock_config_entry_2.add_to_hass(hass)
    mock_config_entries = [mock_config_entry, mock_config_entry_2]

    device_1 = device_registry.async_get_or_create(
        config_entry_id=mock_config_entry.entry_id,
        identifiers={(DOMAIN, mock_config_entry.unique_id)},
        name=mock_config_entry.title,
        entry_type=dr.DeviceEntryType.SERVICE,
    )
    # A stale disabled_by flag can't be set through the registry API, which
    # validates it against the config entry's disabled state; write it
    # directly to simulate existing storage.
    device_1 = attr.evolve(device_1, disabled_by=DeviceEntryDisabler.CONFIG_ENTRY)
    device_registry.devices[device_1.id] = device_1
    entity_registry.async_get_or_create(
        "sensor",
        DOMAIN,
        mock_config_entry.unique_id,
        config_entry=mock_config_entry,
        device_id=device_1.id,
        suggested_object_id="de_jongweg_utrecht_air_quality_index",
        disabled_by=RegistryEntryDisabler.CONFIG_ENTRY,
    )

    device_2 = device_registry.async_get_or_create(
        config_entry_id=mock_config_entry_2.entry_id,
        identifiers={(DOMAIN, mock_config_entry_2.unique_id)},
        name=mock_config_entry_2.title,
        entry_type=dr.DeviceEntryType.SERVICE,
    )
    # A device created on a disabled config entry is disabled by the registry
    # API; clear the flag directly to simulate existing storage with a stale
    # enabled device.
    device_2 = attr.evolve(device_2, disabled_by=None)
    device_registry.devices[device_2.id] = device_2
    entity_registry.async_get_or_create(
        "sensor",
        DOMAIN,
        mock_config_entry_2.unique_id,
        config_entry=mock_config_entry_2,
        device_id=device_2.id,
        suggested_object_id="not_de_jongweg_utrecht_air_quality_index",
    )

    devices = [device_1, device_2]

    # Run migration
    with patch(
        "homeassistant.components.waqi.async_setup_entry",
        return_value=True,
    ):
        await hass.config_entries.async_setup(mock_config_entry.entry_id)
        await hass.async_block_till_done()

    entries = hass.config_entries.async_entries(DOMAIN)
    assert len(entries) == 1
    entry = entries[0]
    assert entry.disabled_by is merged_config_entry_disabled_by
    assert entry.version == 2
    assert entry.minor_version == 1
    assert not entry.options
    assert entry.title == "WAQI"
    assert len(entry.subentries) == 2
    station_subentries = [
        subentry
        for subentry in entry.subentries.values()
        if subentry.subentry_type == "station"
    ]
    assert len(station_subentries) == 2
    for subentry in station_subentries:
        assert subentry.data == {CONF_STATION_NUMBER: int(subentry.unique_id)}
        assert "de Jongweg" in subentry.title

    assert not device_registry.async_get_device(
        identifiers={(DOMAIN, mock_config_entry.entry_id)}
    )
    assert not device_registry.async_get_device(
        identifiers={(DOMAIN, mock_config_entry_2.entry_id)}
    )

    for idx, subentry in enumerate(station_subentries):
        subentry_data = sensor_subentry_data[idx]
        entity = entity_registry.async_get(subentry_data["sensor_entity_id"])
        assert entity.unique_id == subentry.unique_id
        assert entity.config_subentry_id == subentry.subentry_id
        assert entity.config_entry_id == entry.entry_id
        assert entity.disabled_by is subentry_data["entity_disabled_by"]

        assert (
            device := device_registry.async_get_device(
                identifiers={(DOMAIN, subentry.unique_id)}
            )
        )
        assert device.identifiers == {(DOMAIN, subentry.unique_id)}
        assert device.id == devices[subentry_data["device"]].id
        assert device.config_entries == {
            mock_config_entries[main_config_entry].entry_id
        }
        assert device.config_entries_subentries == {
            mock_config_entries[main_config_entry].entry_id: {subentry.subentry_id}
        }
        assert device.disabled_by is subentry_data["device_disabled_by"]