core/tests/components/shelly/__init__.py

"""Tests for the Shelly integration."""

from collections.abc import Mapping, Sequence
from contextlib import contextmanager
from copy import deepcopy
from datetime import timedelta
from typing import Any
from unittest.mock import Mock, patch

from aioshelly.const import MODEL_25
from freezegun.api import FrozenDateTimeFactory
import pytest
from syrupy.assertion import SnapshotAssertion
from syrupy.filters import props

from homeassistant.components.shelly import (
    BLOCK_SLEEPING_PLATFORMS,
    PLATFORMS,
    RPC_SLEEPING_PLATFORMS,
)
from homeassistant.components.shelly.const import (
    CONF_GEN,
    CONF_SLEEP_PERIOD,
    DOMAIN,
    REST_SENSORS_UPDATE_INTERVAL,
    RPC_SENSORS_POLLING_INTERVAL,
)
from homeassistant.config_entries import ConfigEntry
from homeassistant.const import CONF_HOST, CONF_MODEL, Platform
from homeassistant.core import HomeAssistant
from homeassistant.helpers import entity_registry as er
from homeassistant.helpers.device_registry import (
    CONNECTION_NETWORK_MAC,
    DeviceEntry,
    DeviceRegistry,
)

from tests.common import MockConfigEntry, async_fire_time_changed

MOCK_MAC = "123456789ABC"


async def init_integration(
    hass: HomeAssistant,
    gen: int | None,
    model=MODEL_25,
    sleep_period=0,
    options: dict[str, Any] | None = None,
    skip_setup: bool = False,
    data: dict[str, Any] | None = None,
) -> MockConfigEntry:
    """Set up the Shelly integration in Home Assistant."""
    if data is None:
        data = {
            CONF_HOST: "192.168.1.37",
            CONF_SLEEP_PERIOD: sleep_period,
            CONF_MODEL: model,
        }
    if gen is not None:
        data[CONF_GEN] = gen

    entry = MockConfigEntry(
        domain=DOMAIN, data=data, unique_id=MOCK_MAC, options=options, title="Test name"
    )
    entry.add_to_hass(hass)

    if not skip_setup:
        await hass.config_entries.async_setup(entry.entry_id)
        await hass.async_block_till_done()

    return entry


def mutate_rpc_device_status(
    monkeypatch: pytest.MonkeyPatch,
    mock_rpc_device: Mock,
    top_level_key: str,
    key: str,
    value: Any,
) -> None:
    """Mutate status for rpc device."""
    new_status = deepcopy(mock_rpc_device.status)
    new_status[top_level_key][key] = value
    monkeypatch.setattr(mock_rpc_device, "status", new_status)


def inject_rpc_device_event(
    monkeypatch: pytest.MonkeyPatch,
    mock_rpc_device: Mock,
    event: Mapping[str, list[dict[str, Any]] | float],
) -> None:
    """Inject event for rpc device."""
    monkeypatch.setattr(mock_rpc_device, "event", event)
    mock_rpc_device.mock_event()


async def mock_rest_update(
    hass: HomeAssistant,
    freezer: FrozenDateTimeFactory,
    seconds=REST_SENSORS_UPDATE_INTERVAL,
) -> None:
    """Move time to create REST sensors update event."""
    freezer.tick(timedelta(seconds=seconds))
    async_fire_time_changed(hass)
    await hass.async_block_till_done()


async def mock_polling_rpc_update(
    hass: HomeAssistant,
    freezer: FrozenDateTimeFactory,
    seconds: float = RPC_SENSORS_POLLING_INTERVAL,
) -> None:
    """Move time to create polling RPC sensors update event."""
    freezer.tick(timedelta(seconds=seconds))
    async_fire_time_changed(hass)
    await hass.async_block_till_done()


def register_entity(
    hass: HomeAssistant,
    domain: str,
    object_id: str,
    unique_id: str,
    config_entry: ConfigEntry | None = None,
    capabilities: Mapping[str, Any] | None = None,
    device_id: str | None = None,
) -> str:
    """Register enabled entity, return entity_id."""
    entity_registry = er.async_get(hass)
    entity_registry.async_get_or_create(
        domain,
        DOMAIN,
        f"{MOCK_MAC}-{unique_id}",
        suggested_object_id=object_id,
        disabled_by=None,
        config_entry=config_entry,
        capabilities=capabilities,
        device_id=device_id,
    )
    return f"{domain}.{object_id}"


def get_entity(
    hass: HomeAssistant,
    domain: str,
    unique_id: str,
) -> str | None:
    """Get Shelly entity."""
    entity_registry = er.async_get(hass)
    return entity_registry.async_get_entity_id(
        domain, DOMAIN, f"{MOCK_MAC}-{unique_id}"
    )


def register_device(
    device_registry: DeviceRegistry, config_entry: ConfigEntry
) -> DeviceEntry:
    """Register Shelly device."""
    return device_registry.async_get_or_create(
        config_entry_id=config_entry.entry_id,
        connections={(CONNECTION_NETWORK_MAC, MOCK_MAC)},
    )


def register_sub_device(
    device_registry: DeviceRegistry, config_entry: ConfigEntry, unique_id: str
) -> DeviceEntry:
    """Register Shelly sub-device."""
    return device_registry.async_get_or_create(
        config_entry_id=config_entry.entry_id,
        identifiers={(DOMAIN, f"{MOCK_MAC}-{unique_id}")},
        via_device=(DOMAIN, MOCK_MAC),
    )


async def snapshot_device_entities(
    hass: HomeAssistant,
    entity_registry: er.EntityRegistry,
    snapshot: SnapshotAssertion,
    config_entry_id: str,
) -> None:
    """Snapshot all device entities."""

    def sort_event_types(data: Any, path: Sequence[tuple[str, Any]]) -> Any:
        """Sort the event_types list for event entity."""
        if path and path[-1][0] == "event_types" and isinstance(data, list):
            return sorted(data)

        return data

    entity_entries = er.async_entries_for_config_entry(entity_registry, config_entry_id)
    assert entity_entries

    for entity_entry in entity_entries:
        assert entity_entry == snapshot(
            name=f"{entity_entry.entity_id}-entry", exclude=props("event_types")
        )
        assert entity_entry.disabled_by is None, "Please enable all entities."
        state = hass.states.get(entity_entry.entity_id)
        assert state, f"State not found for {entity_entry.entity_id}"
        assert state == snapshot(
            name=f"{entity_entry.entity_id}-state", matcher=sort_event_types
        )


async def force_uptime_value(
    hass: HomeAssistant,
    freezer: FrozenDateTimeFactory,
) -> None:
    """Force time to a specific point."""
    await hass.config.async_set_time_zone("UTC")
    freezer.move_to("2025-05-26 16:04:00+00:00")


@contextmanager
def patch_platforms(platforms: list[Platform]):
    """Only allow given platforms to be loaded."""
    with (
        patch(
            "homeassistant.components.shelly.PLATFORMS",
            list(set(PLATFORMS) & set(platforms)),
        ),
        patch(
            "homeassistant.components.shelly.BLOCK_SLEEPING_PLATFORMS",
            list(set(BLOCK_SLEEPING_PLATFORMS) & set(platforms)),
        ),
        patch(
            "homeassistant.components.shelly.RPC_SLEEPING_PLATFORMS",
            list(set(RPC_SLEEPING_PLATFORMS) & set(platforms)),
        ),
    ):
        yield