core/tests/components/tplink/__init__.py

"""Tests for the TP-Link component."""

from collections import namedtuple
from datetime import datetime
from typing import Any
from unittest.mock import AsyncMock, MagicMock, patch

from kasa import (
    Device,
    DeviceConfig,
    DeviceConnectionParameters,
    DeviceEncryptionType,
    DeviceFamily,
    DeviceType,
    Feature,
    KasaException,
    Module,
)
from kasa.interfaces import Fan, Light, LightEffect, LightState
from kasa.protocol import BaseProtocol
from syrupy import SnapshotAssertion

from homeassistant.components.tplink import (
    CONF_ALIAS,
    CONF_CREDENTIALS_HASH,
    CONF_DEVICE_CONFIG,
    CONF_HOST,
    CONF_MODEL,
    Credentials,
)
from homeassistant.components.tplink.const import DOMAIN
from homeassistant.config_entries import ConfigEntry
from homeassistant.const import Platform
from homeassistant.core import HomeAssistant
from homeassistant.helpers import device_registry as dr, entity_registry as er
from homeassistant.helpers.translation import async_get_translations
from homeassistant.helpers.typing import UNDEFINED
from homeassistant.setup import async_setup_component

from tests.common import MockConfigEntry, load_json_value_fixture

ColorTempRange = namedtuple("ColorTempRange", ["min", "max"])  # noqa: PYI024

MODULE = "homeassistant.components.tplink"
MODULE_CONFIG_FLOW = "homeassistant.components.tplink.config_flow"
IP_ADDRESS = "127.0.0.1"
IP_ADDRESS2 = "127.0.0.2"
ALIAS = "My Bulb"
MODEL = "HS100"
MAC_ADDRESS = "aa:bb:cc:dd:ee:ff"
DEVICE_ID = "123456789ABCDEFGH"
DEVICE_ID_MAC = "AA:BB:CC:DD:EE:FF"
DHCP_FORMATTED_MAC_ADDRESS = MAC_ADDRESS.replace(":", "")
MAC_ADDRESS2 = "11:22:33:44:55:66"
DEFAULT_ENTRY_TITLE = f"{ALIAS} {MODEL}"
CREDENTIALS_HASH_LEGACY = ""
DEVICE_CONFIG_LEGACY = DeviceConfig(IP_ADDRESS)
DEVICE_CONFIG_DICT_LEGACY = DEVICE_CONFIG_LEGACY.to_dict(exclude_credentials=True)
CREDENTIALS = Credentials("foo", "bar")
CREDENTIALS_HASH_AES = "AES/abcdefghijklmnopqrstuvabcdefghijklmnopqrstuv=="
CREDENTIALS_HASH_KLAP = "KLAP/abcdefghijklmnopqrstuv=="
DEVICE_CONFIG_KLAP = DeviceConfig(
    IP_ADDRESS,
    credentials=CREDENTIALS,
    connection_type=DeviceConnectionParameters(
        DeviceFamily.SmartTapoPlug, DeviceEncryptionType.Klap
    ),
    uses_http=True,
)
DEVICE_CONFIG_AES = DeviceConfig(
    IP_ADDRESS2,
    credentials=CREDENTIALS,
    connection_type=DeviceConnectionParameters(
        DeviceFamily.SmartTapoPlug, DeviceEncryptionType.Aes
    ),
    uses_http=True,
)
DEVICE_CONFIG_DICT_KLAP = DEVICE_CONFIG_KLAP.to_dict(exclude_credentials=True)
DEVICE_CONFIG_DICT_AES = DEVICE_CONFIG_AES.to_dict(exclude_credentials=True)

CREATE_ENTRY_DATA_LEGACY = {
    CONF_HOST: IP_ADDRESS,
    CONF_ALIAS: ALIAS,
    CONF_MODEL: MODEL,
    CONF_DEVICE_CONFIG: DEVICE_CONFIG_DICT_LEGACY,
}

CREATE_ENTRY_DATA_KLAP = {
    CONF_HOST: IP_ADDRESS,
    CONF_ALIAS: ALIAS,
    CONF_MODEL: MODEL,
    CONF_CREDENTIALS_HASH: CREDENTIALS_HASH_KLAP,
    CONF_DEVICE_CONFIG: DEVICE_CONFIG_DICT_KLAP,
}
CREATE_ENTRY_DATA_AES = {
    CONF_HOST: IP_ADDRESS2,
    CONF_ALIAS: ALIAS,
    CONF_MODEL: MODEL,
    CONF_CREDENTIALS_HASH: CREDENTIALS_HASH_AES,
    CONF_DEVICE_CONFIG: DEVICE_CONFIG_DICT_AES,
}
CONNECTION_TYPE_KLAP = DeviceConnectionParameters(
    DeviceFamily.SmartTapoPlug, DeviceEncryptionType.Klap
)
CONNECTION_TYPE_KLAP_DICT = CONNECTION_TYPE_KLAP.to_dict()
CONNECTION_TYPE_AES = DeviceConnectionParameters(
    DeviceFamily.SmartTapoPlug, DeviceEncryptionType.Aes
)
CONNECTION_TYPE_AES_DICT = CONNECTION_TYPE_AES.to_dict()


def _load_feature_fixtures():
    fixtures = load_json_value_fixture("features.json", DOMAIN)
    for fixture in fixtures.values():
        if isinstance(fixture["value"], str):
            try:
                time = datetime.strptime(fixture["value"], "%Y-%m-%d %H:%M:%S.%f%z")
                fixture["value"] = time
            except ValueError:
                pass
    return fixtures


FEATURES_FIXTURE = _load_feature_fixtures()


async def setup_platform_for_device(
    hass: HomeAssistant, config_entry: ConfigEntry, platform: Platform, device: Device
):
    """Set up a single tplink platform with a device."""
    config_entry.add_to_hass(hass)

    with (
        patch("homeassistant.components.tplink.PLATFORMS", [platform]),
        _patch_discovery(device=device),
        _patch_connect(device=device),
    ):
        await async_setup_component(hass, DOMAIN, {DOMAIN: {}})
        # Good practice to wait background tasks in tests see PR #112726
        await hass.async_block_till_done(wait_background_tasks=True)


async def snapshot_platform(
    hass: HomeAssistant,
    entity_registry: er.EntityRegistry,
    device_registry: dr.DeviceRegistry,
    snapshot: SnapshotAssertion,
    config_entry_id: str,
) -> None:
    """Snapshot a platform."""
    device_entries = dr.async_entries_for_config_entry(device_registry, config_entry_id)
    assert device_entries
    for device_entry in device_entries:
        assert device_entry == snapshot(
            name=f"{device_entry.name}-entry"
        ), f"device entry snapshot failed for {device_entry.name}"

    entity_entries = er.async_entries_for_config_entry(entity_registry, config_entry_id)
    assert entity_entries
    assert (
        len({entity_entry.domain for entity_entry in entity_entries}) == 1
    ), "Please limit the loaded platforms to 1 platform."

    translations = await async_get_translations(hass, "en", "entity", [DOMAIN])
    for entity_entry in entity_entries:
        if entity_entry.translation_key:
            key = f"component.{DOMAIN}.entity.{entity_entry.domain}.{entity_entry.translation_key}.name"
            assert (
                key in translations
            ), f"No translation for entity {entity_entry.unique_id}, expected {key}"
        assert entity_entry == snapshot(
            name=f"{entity_entry.entity_id}-entry"
        ), f"entity entry snapshot failed for {entity_entry.entity_id}"
        if entity_entry.disabled_by is None:
            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"
            ), f"state snapshot failed for {entity_entry.entity_id}"


def _mock_protocol() -> BaseProtocol:
    protocol = MagicMock(spec=BaseProtocol)
    protocol.close = AsyncMock()
    return protocol


def _mocked_device(
    device_config=DEVICE_CONFIG_LEGACY,
    credentials_hash=CREDENTIALS_HASH_LEGACY,
    mac=MAC_ADDRESS,
    device_id=DEVICE_ID,
    alias=ALIAS,
    model=MODEL,
    ip_address: str | None = None,
    modules: list[str] | None = None,
    children: list[Device] | None = None,
    features: list[str | Feature] | None = None,
    device_type=None,
    spec: type = Device,
) -> Device:
    device = MagicMock(spec=spec, name="Mocked device")
    device.update = AsyncMock()
    device.turn_off = AsyncMock()
    device.turn_on = AsyncMock()

    device.mac = mac
    device.alias = alias
    device.model = model
    device.device_id = device_id
    device.hw_info = {"sw_ver": "1.0.0", "hw_ver": "1.0.0"}
    device.modules = {}
    device.features = {}

    if not ip_address:
        ip_address = IP_ADDRESS
    else:
        device_config.host = ip_address
    device.host = ip_address

    if modules:
        device.modules = {
            module_name: MODULE_TO_MOCK_GEN[module_name](device)
            for module_name in modules
        }

    if features:
        device.features = {
            feature_id: _mocked_feature(feature_id, require_fixture=True)
            for feature_id in features
            if isinstance(feature_id, str)
        }

        device.features.update(
            {
                feature.id: feature
                for feature in features
                if isinstance(feature, Feature)
            }
        )
    device.children = []
    if children:
        for child in children:
            child.mac = mac
        device.children = children
    device.device_type = device_type if device_type else DeviceType.Unknown
    if (
        not device_type
        and device.children
        and all(
            child.device_type is DeviceType.StripSocket for child in device.children
        )
    ):
        device.device_type = DeviceType.Strip

    device.protocol = _mock_protocol()
    device.config = device_config
    device.credentials_hash = credentials_hash
    return device


def _mocked_feature(
    id: str,
    *,
    require_fixture=False,
    value: Any = UNDEFINED,
    name=None,
    type_=None,
    category=None,
    precision_hint=None,
    choices=None,
    unit=None,
    minimum_value=0,
    maximum_value=2**16,  # Arbitrary max
) -> Feature:
    """Get a mocked feature.

    If kwargs are provided they will override the attributes for any features defined in fixtures.json
    """
    feature = MagicMock(spec=Feature, name=f"Mocked {id} feature")
    feature.id = id
    feature.name = name or id.upper()
    feature.set_value = AsyncMock()
    if not (fixture := FEATURES_FIXTURE.get(id)):
        assert (
            require_fixture is False
        ), f"No fixture defined for feature {id} and require_fixture is True"
        assert (
            value is not UNDEFINED
        ), f"Value must be provided if feature {id} not defined in features.json"
        fixture = {"value": value, "category": "Primary", "type": "Sensor"}
    elif value is not UNDEFINED:
        fixture["value"] = value
    feature.value = fixture["value"]

    feature.type = type_ or Feature.Type[fixture["type"]]
    feature.category = category or Feature.Category[fixture["category"]]

    # sensor
    feature.precision_hint = precision_hint or fixture.get("precision_hint")
    feature.unit = unit or fixture.get("unit")

    # number
    feature.minimum_value = minimum_value or fixture.get("minimum_value")
    feature.maximum_value = maximum_value or fixture.get("maximum_value")

    # select
    feature.choices = choices or fixture.get("choices")
    return feature


def _mocked_light_module(device) -> Light:
    light = MagicMock(spec=Light, name="Mocked light module")
    light.update = AsyncMock()
    light.brightness = 50
    light.color_temp = 4000
    light.state = LightState(
        light_on=True, brightness=light.brightness, color_temp=light.color_temp
    )
    light.is_color = True
    light.is_variable_color_temp = True
    light.is_dimmable = True
    light.is_brightness = True
    light.has_effects = False
    light.hsv = (10, 30, 5)
    light.valid_temperature_range = ColorTempRange(min=4000, max=9000)
    light.hw_info = {"sw_ver": "1.0.0", "hw_ver": "1.0.0"}

    async def _set_state(state, *_, **__):
        light.state = state

    light.set_state = AsyncMock(wraps=_set_state)

    async def _set_brightness(brightness, *_, **__):
        light.state.brightness = brightness
        light.state.light_on = brightness > 0

    light.set_brightness = AsyncMock(wraps=_set_brightness)

    async def _set_hsv(h, s, v, *_, **__):
        light.state.hue = h
        light.state.saturation = s
        light.state.brightness = v
        light.state.light_on = True

    light.set_hsv = AsyncMock(wraps=_set_hsv)

    async def _set_color_temp(temp, *_, **__):
        light.state.color_temp = temp
        light.state.light_on = True

    light.set_color_temp = AsyncMock(wraps=_set_color_temp)
    light.protocol = _mock_protocol()
    return light


def _mocked_light_effect_module(device) -> LightEffect:
    effect = MagicMock(spec=LightEffect, name="Mocked light effect")
    effect.has_effects = True
    effect.has_custom_effects = True
    effect.effect = "Effect1"
    effect.effect_list = ["Off", "Effect1", "Effect2"]

    async def _set_effect(effect_name, *_, **__):
        assert (
            effect_name in effect.effect_list
        ), f"set_effect '{effect_name}' not in {effect.effect_list}"
        assert device.modules[
            Module.Light
        ], "Need a light module to test set_effect method"
        device.modules[Module.Light].state.light_on = True
        effect.effect = effect_name

    effect.set_effect = AsyncMock(wraps=_set_effect)
    effect.set_custom_effect = AsyncMock()
    return effect


def _mocked_fan_module(effect) -> Fan:
    fan = MagicMock(auto_spec=Fan, name="Mocked fan")
    fan.fan_speed_level = 0
    fan.set_fan_speed_level = AsyncMock()
    return fan


def _mocked_strip_children(features=None, alias=None) -> list[Device]:
    plug0 = _mocked_device(
        alias="Plug0" if alias is None else alias,
        device_id="bb:bb:cc:dd:ee:ff_PLUG0DEVICEID",
        mac="bb:bb:cc:dd:ee:ff",
        device_type=DeviceType.StripSocket,
        features=features,
    )
    plug1 = _mocked_device(
        alias="Plug1" if alias is None else alias,
        device_id="cc:bb:cc:dd:ee:ff_PLUG1DEVICEID",
        mac="cc:bb:cc:dd:ee:ff",
        device_type=DeviceType.StripSocket,
        features=features,
    )
    plug0.is_on = True
    plug1.is_on = False
    return [plug0, plug1]


def _mocked_energy_features(
    power=None, total=None, voltage=None, current=None, today=None
) -> list[Feature]:
    feats = []
    if power is not None:
        feats.append(
            _mocked_feature(
                "current_consumption",
                value=power,
            )
        )
    if total is not None:
        feats.append(
            _mocked_feature(
                "consumption_total",
                value=total,
            )
        )
    if voltage is not None:
        feats.append(
            _mocked_feature(
                "voltage",
                value=voltage,
            )
        )
    if current is not None:
        feats.append(
            _mocked_feature(
                "current",
                value=current,
            )
        )
    # Today is always reported as 0 by the library rather than none
    feats.append(
        _mocked_feature(
            "consumption_today",
            value=today if today is not None else 0.0,
        )
    )
    return feats


MODULE_TO_MOCK_GEN = {
    Module.Light: _mocked_light_module,
    Module.LightEffect: _mocked_light_effect_module,
    Module.Fan: _mocked_fan_module,
}


def _patch_discovery(device=None, no_device=False):
    async def _discovery(*args, **kwargs):
        if no_device:
            return {}
        return {IP_ADDRESS: _mocked_device()}

    return patch("homeassistant.components.tplink.Discover.discover", new=_discovery)


def _patch_single_discovery(device=None, no_device=False):
    async def _discover_single(*args, **kwargs):
        if no_device:
            raise KasaException
        return device if device else _mocked_device()

    return patch(
        "homeassistant.components.tplink.Discover.discover_single", new=_discover_single
    )


def _patch_connect(device=None, no_device=False):
    async def _connect(*args, **kwargs):
        if no_device:
            raise KasaException
        return device if device else _mocked_device()

    return patch("homeassistant.components.tplink.Device.connect", new=_connect)


async def initialize_config_entry_for_device(
    hass: HomeAssistant, dev: Device
) -> MockConfigEntry:
    """Create a mocked configuration entry for the given device.

    Note, the rest of the tests should probably be converted over to use this
    instead of repeating the initialization routine for each test separately
    """
    config_entry = MockConfigEntry(
        title="TP-Link", domain=DOMAIN, unique_id=dev.mac, data={CONF_HOST: dev.host}
    )
    config_entry.add_to_hass(hass)

    with (
        _patch_discovery(device=dev),
        _patch_single_discovery(device=dev),
        _patch_connect(device=dev),
    ):
        await hass.config_entries.async_setup(config_entry.entry_id)
        await hass.async_block_till_done()

    return config_entry