test_number.py

"""Tests for the WLED number platform."""

from collections.abc import Generator
from unittest.mock import MagicMock, patch

from freezegun.api import FrozenDateTimeFactory
import pytest
from syrupy.assertion import SnapshotAssertion
from wled import Device as WLEDDevice, WLEDConnectionError, WLEDError

from homeassistant.components.number import (
    ATTR_VALUE,
    DOMAIN as NUMBER_DOMAIN,
    SERVICE_SET_VALUE,
)
from homeassistant.components.wled.const import (
    CONF_KEEP_MAIN_LIGHT,
    DOMAIN,
    SCAN_INTERVAL,
)
from homeassistant.const import ATTR_ENTITY_ID, STATE_UNAVAILABLE, Platform
from homeassistant.core import HomeAssistant
from homeassistant.exceptions import HomeAssistantError
from homeassistant.helpers import device_registry as dr, entity_registry as er

from tests.common import (
    MockConfigEntry,
    async_fire_time_changed,
    async_load_json_object_fixture,
    snapshot_platform,
)

pytestmark = pytest.mark.usefixtures("init_integration")


@pytest.fixture(autouse=True)
def override_platforms() -> Generator[None]:
    """Override PLATFORMS."""
    with patch("homeassistant.components.wled.PLATFORMS", [Platform.NUMBER]):
        yield


async def test_snapshots(
    hass: HomeAssistant,
    entity_registry: er.EntityRegistry,
    snapshot: SnapshotAssertion,
    mock_config_entry: MockConfigEntry,
) -> None:
    """Test snapshot of the platform."""
    await snapshot_platform(hass, entity_registry, snapshot, mock_config_entry.entry_id)


@pytest.mark.parametrize(
    ("entity_id", "value", "called_arg"),
    [
        ("number.wled_rgb_light_segment_1_speed", 42, "speed"),
        ("number.wled_rgb_light_segment_1_intensity", 42, "intensity"),
    ],
)
async def test_numbers(
    hass: HomeAssistant,
    device_registry: dr.DeviceRegistry,
    entity_registry: er.EntityRegistry,
    snapshot: SnapshotAssertion,
    mock_wled: MagicMock,
    entity_id: str,
    value: int,
    called_arg: str,
) -> None:
    """Test the creation and values of the WLED numbers."""
    assert (state := hass.states.get(entity_id))
    assert state == snapshot

    assert (entity_entry := entity_registry.async_get(state.entity_id))
    assert entity_entry == snapshot

    assert entity_entry.device_id
    assert (device_entry := device_registry.async_get(entity_entry.device_id))
    assert device_entry == snapshot

    # Test a regular state change service call
    await hass.services.async_call(
        NUMBER_DOMAIN,
        SERVICE_SET_VALUE,
        {ATTR_ENTITY_ID: entity_id, ATTR_VALUE: value},
        blocking=True,
    )

    assert mock_wled.segment.call_count == 1
    mock_wled.segment.assert_called_with(segment_id=1, **{called_arg: value})

    # Test with WLED error
    mock_wled.segment.side_effect = WLEDError
    with pytest.raises(HomeAssistantError, match="Invalid response from WLED API"):
        await hass.services.async_call(
            NUMBER_DOMAIN,
            SERVICE_SET_VALUE,
            {ATTR_ENTITY_ID: entity_id, ATTR_VALUE: value},
            blocking=True,
        )
    assert mock_wled.segment.call_count == 2

    # Ensure the entity is still available
    assert (state := hass.states.get(entity_id))
    assert state.state != STATE_UNAVAILABLE

    # Test when a connection error occurs
    mock_wled.segment.side_effect = WLEDConnectionError
    with pytest.raises(HomeAssistantError, match="Error communicating with WLED API"):
        await hass.services.async_call(
            NUMBER_DOMAIN,
            SERVICE_SET_VALUE,
            {ATTR_ENTITY_ID: entity_id, ATTR_VALUE: value},
            blocking=True,
        )
    assert mock_wled.segment.call_count == 3

    # Ensure the entity became unavailable after the connection error
    assert (state := hass.states.get(entity_id))
    assert state.state == STATE_UNAVAILABLE


@pytest.mark.parametrize("device_fixture", ["rgb_single_segment"])
@pytest.mark.parametrize(
    ("entity_id_segment0", "state_segment0", "entity_id_segment1", "state_segment1"),
    [
        (
            "number.wled_rgb_light_speed",
            "32",
            "number.wled_rgb_light_segment_1_speed",
            "16",
        ),
        (
            "number.wled_rgb_light_intensity",
            "128",
            "number.wled_rgb_light_segment_1_intensity",
            "64",
        ),
    ],
)
async def test_speed_dynamically_handle_segments(
    hass: HomeAssistant,
    freezer: FrozenDateTimeFactory,
    mock_wled: MagicMock,
    entity_id_segment0: str,
    entity_id_segment1: str,
    state_segment0: str,
    state_segment1: str,
) -> None:
    """Test if a new/deleted segment is dynamically added/removed."""
    assert (segment0 := hass.states.get(entity_id_segment0))
    assert segment0.state == state_segment0
    assert not hass.states.get(entity_id_segment1)

    # Test adding a segment dynamically...
    return_value = mock_wled.update.return_value
    mock_wled.update.return_value = WLEDDevice.from_dict(
        await async_load_json_object_fixture(hass, "rgb.json", DOMAIN)
    )

    freezer.tick(SCAN_INTERVAL)
    async_fire_time_changed(hass)
    await hass.async_block_till_done()

    assert (segment0 := hass.states.get(entity_id_segment0))
    assert segment0.state == state_segment0
    assert (segment1 := hass.states.get(entity_id_segment1))
    assert segment1.state == state_segment1

    # Test remove segment again...
    mock_wled.update.return_value = return_value
    freezer.tick(SCAN_INTERVAL)
    async_fire_time_changed(hass)
    await hass.async_block_till_done()

    assert (segment0 := hass.states.get(entity_id_segment0))
    assert segment0.state == state_segment0
    assert (segment1 := hass.states.get(entity_id_segment1))
    assert segment1.state == STATE_UNAVAILABLE


@pytest.mark.parametrize(
    ("device_fixture", "keep_main_light", "expected_entities"),
    [
        (
            "rgb_single_segment",
            False,
            {
                "number.wled_rgb_light_speed": "WLED RGB Light Speed",
                "number.wled_rgb_light_intensity": "WLED RGB Light Intensity",
            },
        ),
        (
            "rgb_single_segment",
            True,
            {
                "number.wled_rgb_light_speed": "WLED RGB Light Segment 0 speed",
                "number.wled_rgb_light_intensity": "WLED RGB Light Segment 0 intensity",
            },
        ),
        (
            "rgb",
            False,
            {
                "number.wled_rgb_light_speed": "WLED RGB Light Speed",
                "number.wled_rgb_light_intensity": "WLED RGB Light Intensity",
                "number.wled_rgb_light_segment_1_speed": "WLED RGB Light Segment 1 speed",
                "number.wled_rgb_light_segment_1_intensity": "WLED RGB Light Segment 1 intensity",
            },
        ),
        (
            "rgb",
            True,
            {
                "number.wled_rgb_light_speed": "WLED RGB Light Segment 0 speed",
                "number.wled_rgb_light_intensity": "WLED RGB Light Segment 0 intensity",
                "number.wled_rgb_light_segment_1_speed": "WLED RGB Light Segment 1 speed",
                "number.wled_rgb_light_segment_1_intensity": "WLED RGB Light Segment 1 intensity",
            },
        ),
    ],
)
async def test_keep_main_light_entity_names(
    hass: HomeAssistant,
    init_integration: MockConfigEntry,
    keep_main_light: bool,
    expected_entities: dict[str, str],
) -> None:
    """Test entity friendly names with and without the keep_main_light option."""
    hass.config_entries.async_update_entry(
        init_integration, options={CONF_KEEP_MAIN_LIGHT: keep_main_light}
    )
    await hass.config_entries.async_reload(init_integration.entry_id)
    await hass.async_block_till_done()

    actual_entities = {
        state.entity_id: state.attributes.get("friendly_name")
        for state in hass.states.async_all("number")
    }
    assert actual_entities == expected_entities