ha-ninlight/tests/test_update.py

"""Firmware update entity: offer, Wi-Fi guard, install flow."""

from __future__ import annotations

import asyncio
from copy import deepcopy
from pathlib import Path
from typing import Any
from unittest.mock import MagicMock
from urllib.parse import urlsplit

from homeassistant.components.update import (
    ATTR_IN_PROGRESS,
    ATTR_INSTALLED_VERSION,
    ATTR_LATEST_VERSION,
    SERVICE_INSTALL,
)
from homeassistant.components.update import (
    DOMAIN as UPDATE_DOMAIN,
)
from homeassistant.const import ATTR_ENTITY_ID, CONF_HOST, CONF_TOKEN, STATE_OFF, STATE_ON
from homeassistant.core import HomeAssistant
from homeassistant.core_config import async_process_ha_core_config
from homeassistant.exceptions import HomeAssistantError
from homeassistant.helpers import issue_registry as ir
import pytest
from pytest_homeassistant_custom_component.common import MockConfigEntry
from pytest_homeassistant_custom_component.typing import ClientSessionGenerator

from custom_components.ninlight import update as update_module
from custom_components.ninlight.api import NinlightBusyError, NinlightConnectionError
from custom_components.ninlight.const import CONF_LAMP_BASE_URL, DOMAIN
from custom_components.ninlight.shared import DATA_SHARED, async_scan_catalog
from custom_components.ninlight.update import depends_on_compiled_wifi
from tests.helpers import entity_id_for
from tests_unit.images import CONTENT_WITH_SSID, CONTENT_WITHOUT_SSID, build_image
from tests_unit.samples import HOST_RGB, MAC_RGB, SETTINGS_COMPILED, SETTINGS_STORED, TOKEN

BASE = "http://192.168.1.50:8123"
INSTALLED = "0.3.8-36-gca9bb"
NEWER = "0.3.8-47-g71a57"
NEWEST = "0.3.9-rc1"
EVEN_NEWER = "0.3.9-rc2"  # newer than NEWEST


@pytest.fixture
def mock_config_entry() -> MockConfigEntry:
    return MockConfigEntry(
        domain=DOMAIN,
        title="Klyqa GU10 RGB CW/WW 06b8",
        unique_id=MAC_RGB,
        data={CONF_HOST: HOST_RGB, CONF_TOKEN: TOKEN},
        options={CONF_LAMP_BASE_URL: BASE},
    )


@pytest.fixture(autouse=True)
def fast_polling(monkeypatch: pytest.MonkeyPatch) -> None:
    monkeypatch.setattr(update_module, "OTA_POLL_INTERVAL_S", 0)


def put(directory: Path, name: str, **kwargs: Any) -> Path:
    path = directory / name
    path.write_bytes(build_image(**kwargs))
    return path


async def rescan(hass: HomeAssistant) -> None:
    await async_scan_catalog(hass, hass.data[DATA_SHARED])
    await hass.async_block_till_done()


def wifi_issue(hass: HomeAssistant, entry: MockConfigEntry) -> ir.IssueEntry | None:
    return ir.async_get(hass).async_get_issue(DOMAIN, f"wifi_guard_{entry.entry_id}")


async def install(hass: HomeAssistant) -> None:
    await hass.services.async_call(
        UPDATE_DOMAIN,
        SERVICE_INSTALL,
        {ATTR_ENTITY_ID: entity_id_for(hass, "update", "firmware")},
        blocking=True,
    )


def test_only_stored_wifi_is_independent() -> None:
    assert not depends_on_compiled_wifi(SETTINGS_STORED)
    assert depends_on_compiled_wifi(SETTINGS_COMPILED)
    assert depends_on_compiled_wifi({"wifi": {"ssid": ""}})
    assert depends_on_compiled_wifi({"wifi": {"source": "flash"}})
    assert depends_on_compiled_wifi({})


async def test_no_images_no_update(hass: HomeAssistant, init_integration) -> None:
    state = hass.states.get(entity_id_for(hass, "update", "firmware"))
    assert state.state == STATE_OFF
    assert state.attributes[ATTR_INSTALLED_VERSION] == INSTALLED
    assert state.attributes[ATTR_LATEST_VERSION] == INSTALLED


async def test_newer_image_with_credentials_is_offered(
    hass: HomeAssistant, init_integration, firmware_dir: Path
) -> None:
    put(firmware_dir, "a.img", version=NEWER, content=CONTENT_WITH_SSID)
    await rescan(hass)
    state = hass.states.get(entity_id_for(hass, "update", "firmware"))
    assert state.state == STATE_ON
    assert state.attributes[ATTR_LATEST_VERSION] == NEWER
    assert wifi_issue(hass, init_integration) is None


@pytest.mark.parametrize(
    "image",
    [
        {
            "device_type": "gu10_cw",
            "version": "0.4.0",
        },  # other type that only differs in the suffix
        {"device_type": "gu10_rgb", "version": "0.3.8-36-gffff0"},  # same version, other hash
        {"device_type": "gu10_rgb", "version": "0.3.8"},  # older
    ],
)
async def test_unsuitable_images_are_not_offered(
    hass: HomeAssistant, init_integration, firmware_dir: Path, image: dict[str, str]
) -> None:
    put(firmware_dir, "a.img", content=CONTENT_WITH_SSID, **image)
    await rescan(hass)
    assert hass.states.get(entity_id_for(hass, "update", "firmware")).state == STATE_OFF


async def test_wifi_guard_blocks_image_without_credentials(
    hass: HomeAssistant, init_integration, firmware_dir: Path
) -> None:
    nocreds = put(firmware_dir, "nocreds.img", version=NEWEST, content=CONTENT_WITHOUT_SSID)
    await rescan(hass)
    entity_id = entity_id_for(hass, "update", "firmware")
    assert hass.states.get(entity_id).state == STATE_OFF
    issue = wifi_issue(hass, init_integration)
    assert issue is not None
    assert issue.translation_placeholders == {
        "name": "Klyqa GU10 RGB CW/WW 06b8",
        "version": NEWEST,
        "device_type": "gu10_rgb",
        "ssid": "ninIOT",
    }

    put(firmware_dir, "creds.img", version=NEWER, content=CONTENT_WITH_SSID)
    await rescan(hass)
    assert hass.states.get(entity_id).attributes[ATTR_LATEST_VERSION] == NEWER
    assert wifi_issue(hass, init_integration) is not None  # the newer, blocked one is still there

    nocreds.unlink()
    await rescan(hass)
    assert hass.states.get(entity_id).attributes[ATTR_LATEST_VERSION] == NEWER
    assert wifi_issue(hass, init_integration) is None


async def test_lamp_with_stored_wifi_gets_image_without_credentials(
    hass: HomeAssistant, init_integration, firmware_dir: Path, lamp_settings: dict[str, Any]
) -> None:
    lamp_settings.clear()
    lamp_settings.update(SETTINGS_STORED)
    put(firmware_dir, "nocreds.img", version=NEWEST, content=CONTENT_WITHOUT_SSID)
    await rescan(hass)
    assert (
        hass.states.get(entity_id_for(hass, "update", "firmware")).attributes[ATTR_LATEST_VERSION]
        == NEWEST
    )
    assert wifi_issue(hass, init_integration) is None


async def test_lamp_without_ssid_gets_nothing(
    hass: HomeAssistant, init_integration, firmware_dir: Path, lamp_info: dict[str, Any]
) -> None:
    lamp_info["wifi"]["ssid"] = ""
    await init_integration.runtime_data.async_refresh()
    put(firmware_dir, "creds.img", version=NEWER, content=CONTENT_WITH_SSID)
    await rescan(hass)
    assert hass.states.get(entity_id_for(hass, "update", "firmware")).state == STATE_OFF


async def test_unreadable_settings_keep_previous_offer(
    hass: HomeAssistant, init_integration, firmware_dir: Path, mock_client: MagicMock
) -> None:
    mock_client.get_settings.side_effect = NinlightConnectionError("down")
    put(firmware_dir, "creds.img", version=NEWER, content=CONTENT_WITH_SSID)
    await rescan(hass)
    assert hass.states.get(entity_id_for(hass, "update", "firmware")).state == STATE_OFF


async def test_install_serves_image_once_and_waits_for_new_version(
    hass: HomeAssistant,
    hass_client_no_auth: ClientSessionGenerator,
    init_integration,
    mock_client: MagicMock,
    lamp_info: dict[str, Any],
    firmware_dir: Path,
) -> None:
    path = put(firmware_dir, "a.img", version=NEWER, content=CONTENT_WITH_SSID)
    await rescan(hass)
    client = await hass_client_no_auth()
    served: dict[str, Any] = {}

    async def lamp_downloads(url: str) -> None:
        served["url"] = url
        response = await client.get(urlsplit(url).path)
        served["status"] = response.status
        served["body"] = await response.read()
        served["ota_flag"] = init_integration.runtime_data.ota_in_progress
        lamp_info["fw_version"] = NEWER
        lamp_info["uptime"] = 4

    mock_client.fw_update.side_effect = lamp_downloads
    await install(hass)
    await hass.async_block_till_done()

    assert served["url"].startswith(f"{BASE}/api/ninlight/firmware/")
    assert (served["status"], served["body"]) == (200, path.read_bytes())
    assert served["ota_flag"] is True
    assert (await client.get(urlsplit(served["url"]).path)).status == 404
    assert len(hass.data[DATA_SHARED].registry) == 0
    state = hass.states.get(entity_id_for(hass, "update", "firmware"))
    assert state.attributes[ATTR_INSTALLED_VERSION] == NEWER
    assert state.attributes[ATTR_IN_PROGRESS] is False
    assert state.state == STATE_OFF
    assert init_integration.runtime_data.ota_in_progress is False


async def test_install_serves_exactly_the_verified_bytes(
    hass: HomeAssistant,
    hass_client_no_auth: ClientSessionGenerator,
    init_integration,
    mock_client: MagicMock,
    lamp_info: dict[str, Any],
    firmware_dir: Path,
) -> None:
    """F1: the lamp must receive exactly the bytes checked (Wi-Fi guard, type, hash) at install
    time, even if the file on disk is replaced by something else afterwards - e.g. a build script
    that writes the no-credentials variant over the same filename mid-update.
    """
    path = put(firmware_dir, "a.img", version=NEWER, content=CONTENT_WITH_SSID)
    verified_bytes = path.read_bytes()
    await rescan(hass)
    client = await hass_client_no_auth()
    served: dict[str, Any] = {}

    async def lamp_downloads(url: str) -> None:
        # Swap in a different, otherwise-valid image for the same device type right before the
        # lamp "fetches" it. The Wi-Fi guard already passed for `verified_bytes`; this replacement
        # has no credentials at all and must never reach the lamp.
        path.write_bytes(build_image(version=NEWER, content=CONTENT_WITHOUT_SSID))
        served["url"] = url
        response = await client.get(urlsplit(url).path)
        served["status"] = response.status
        served["body"] = await response.read()
        lamp_info["fw_version"] = NEWER
        lamp_info["uptime"] = 4

    mock_client.fw_update.side_effect = lamp_downloads
    await install(hass)
    await hass.async_block_till_done()

    assert served["status"] == 200
    assert served["body"] == verified_bytes
    assert served["body"] != path.read_bytes()  # the swapped-in file, not what was served


async def test_rescan_during_install_does_not_trigger_evaluation(
    hass: HomeAssistant,
    init_integration,
    mock_client: MagicMock,
    lamp_info: dict[str, Any],
    firmware_dir: Path,
) -> None:
    """M2: a catalog rescan landing while an install is in progress must not evaluate offers
    against the (currently flashing) lamp - that would mean extra requests to it and could churn
    the offer/issue underneath the install.
    """
    put(firmware_dir, "a.img", version=NEWER, content=CONTENT_WITH_SSID)
    await rescan(hass)
    calls_before_rescan: dict[str, int] = {}

    async def lamp_flashes(url: str) -> None:
        # a periodic catalog scan happens to land while this install is in progress
        calls_before_rescan["count"] = mock_client.get_settings.await_count
        await rescan(hass)
        lamp_info["fw_version"] = NEWER
        lamp_info["uptime"] = 4

    mock_client.fw_update.side_effect = lamp_flashes
    await install(hass)
    await hass.async_block_till_done()

    assert mock_client.get_settings.await_count == calls_before_rescan["count"]


@pytest.mark.parametrize(
    ("lamp_reaction", "key", "reason_part"),
    [
        ({"last_ota_error": "hash"}, "update_failed", "hash"),
        ({"uptime": 3}, "update_failed", INSTALLED),  # rebooted, still the old firmware
    ],
)
async def test_install_reports_lamp_failures(
    hass: HomeAssistant,
    init_integration,
    mock_client: MagicMock,
    lamp_info: dict[str, Any],
    firmware_dir: Path,
    lamp_reaction: dict[str, Any],
    key: str,
    reason_part: str,
) -> None:
    put(firmware_dir, "a.img", version=NEWER, content=CONTENT_WITH_SSID)
    await rescan(hass)

    async def lamp_fails(url: str) -> None:
        lamp_info.update(lamp_reaction)

    mock_client.fw_update.side_effect = lamp_fails
    with pytest.raises(HomeAssistantError) as raised:
        await install(hass)
    assert raised.value.translation_key == key
    assert reason_part in raised.value.translation_placeholders["reason"]
    assert len(hass.data[DATA_SHARED].registry) == 0
    assert init_integration.runtime_data.ota_in_progress is False


async def test_install_busy_lamp(
    hass: HomeAssistant, init_integration, mock_client: MagicMock, firmware_dir: Path
) -> None:
    put(firmware_dir, "a.img", version=NEWER, content=CONTENT_WITH_SSID)
    await rescan(hass)
    mock_client.fw_update.side_effect = NinlightBusyError(409, "an update is already running")
    with pytest.raises(HomeAssistantError) as raised:
        await install(hass)
    assert raised.value.translation_key == "update_busy"
    assert len(hass.data[DATA_SHARED].registry) == 0


async def test_install_times_out(
    hass: HomeAssistant, init_integration, firmware_dir: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
    monkeypatch.setattr(update_module, "OTA_TIMEOUT_S", 0)
    put(firmware_dir, "a.img", version=NEWER, content=CONTENT_WITH_SSID)
    await rescan(hass)
    with pytest.raises(HomeAssistantError) as raised:
        await install(hass)
    assert raised.value.translation_key == "update_timeout"


async def test_install_needs_http_address(
    hass: HomeAssistant, mock_client: MagicMock, firmware_dir: Path
) -> None:
    await async_process_ha_core_config(hass, {"internal_url": "https://ha.local:8123"})
    entry = MockConfigEntry(
        domain=DOMAIN, unique_id=MAC_RGB, data={CONF_HOST: HOST_RGB, CONF_TOKEN: TOKEN}
    )
    entry.add_to_hass(hass)
    assert await hass.config_entries.async_setup(entry.entry_id)
    put(firmware_dir, "a.img", version=NEWER, content=CONTENT_WITH_SSID)
    await rescan(hass)
    with pytest.raises(HomeAssistantError) as raised:
        await install(hass)
    assert raised.value.translation_key == "no_http_url"
    mock_client.fw_update.assert_not_awaited()


async def test_install_rechecks_wifi_right_before_sending(
    hass: HomeAssistant,
    init_integration,
    mock_client: MagicMock,
    firmware_dir: Path,
    lamp_settings: dict[str, Any],
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    # Guards against a regression that would let this install slip past the re-check into the
    # wait loop: without this it would hang for the full OTA_TIMEOUT_S (12 min) instead of failing
    # fast with "wifi_guard".
    monkeypatch.setattr(update_module, "OTA_TIMEOUT_S", 0)
    lamp_settings.clear()
    lamp_settings.update(SETTINGS_STORED)
    put(firmware_dir, "nocreds.img", version=NEWEST, content=CONTENT_WITHOUT_SSID)
    await rescan(hass)
    assert hass.states.get(entity_id_for(hass, "update", "firmware")).state == STATE_ON

    lamp_settings["wifi"]["source"] = "compiled"  # the lamp forgot its stored Wi-Fi meanwhile
    with pytest.raises(HomeAssistantError) as raised:
        await install(hass)
    assert raised.value.translation_key == "wifi_guard"
    assert raised.value.translation_placeholders == {"ssid": "ninIOT"}
    mock_client.fw_update.assert_not_awaited()


async def test_install_rejects_changed_file(
    hass: HomeAssistant, init_integration, mock_client: MagicMock, firmware_dir: Path
) -> None:
    path = put(firmware_dir, "a.img", version=NEWER, content=CONTENT_WITH_SSID)
    await rescan(hass)
    path.write_bytes(build_image(version=NEWEST, content=CONTENT_WITH_SSID))
    with pytest.raises(HomeAssistantError) as raised:
        await install(hass)
    assert raised.value.translation_key == "image_invalid"
    mock_client.fw_update.assert_not_awaited()


async def test_install_rejects_downgrade(
    hass: HomeAssistant,
    init_integration,
    mock_client: MagicMock,
    lamp_info: dict[str, Any],
    firmware_dir: Path,
) -> None:
    """The lamp may have been reflashed by another route since the offer was computed."""
    put(firmware_dir, "a.img", version=NEWER, content=CONTENT_WITH_SSID)
    await rescan(hass)
    entity_id = entity_id_for(hass, "update", "firmware")
    assert hass.states.get(entity_id).attributes[ATTR_LATEST_VERSION] == NEWER

    # A newer image also appears on disk, and the lamp is reflashed by another route to a version
    # newer than our stale offer - but the follow-up evaluation (triggered by the catalog scan)
    # fails at get_settings and returns early, so the stale NEWER offer survives untouched.
    put(firmware_dir, "b.img", version=EVEN_NEWER, content=CONTENT_WITH_SSID)
    mock_client.get_settings.side_effect = NinlightConnectionError("down")
    lamp_info["fw_version"] = NEWEST
    await rescan(hass)
    assert hass.states.get(entity_id).attributes[ATTR_LATEST_VERSION] == NEWER

    with pytest.raises(HomeAssistantError) as raised:
        await install(hass)
    assert raised.value.translation_key == "no_update"
    mock_client.fw_update.assert_not_awaited()


async def test_install_rejects_concurrent_calls(
    hass: HomeAssistant,
    init_integration,
    mock_client: MagicMock,
    lamp_info: dict[str, Any],
    firmware_dir: Path,
) -> None:
    """A second call arriving in the pre-flag window (before in_progress is set) must be
    rejected by our own lock rather than allowed to touch OTA state.

    HA's own service-level "already in progress" guard only helps once `_attr_in_progress` is
    already True (it is a `cached_property` that HA invalidates whenever `_attr_in_progress` is
    set, so it is never stale) - by the time that guard would fire, a second concurrent call is
    already blocked by HA itself with translation_key "update_in_progress", not "update_busy".
    The real gap I2 describes is *before* that: image verification, `get_info` and `get_settings`
    are all awaited before we ever set the flags, and HA's guard does not protect that window. This
    test
    blocks the first `get_info()` re-check inside that window to reach it deterministically.
    """
    put(firmware_dir, "a.img", version=NEWER, content=CONTENT_WITH_SSID)
    await rescan(hass)
    entity_id = entity_id_for(hass, "update", "firmware")

    entered = asyncio.Event()
    release = asyncio.Event()
    first_call = True

    async def blocking_get_info() -> dict[str, Any]:
        nonlocal first_call
        if first_call:
            first_call = False
            entered.set()
            await release.wait()
        return deepcopy(lamp_info)

    async def lamp_flashes(url: str) -> None:
        lamp_info["fw_version"] = NEWER
        lamp_info["uptime"] = 4

    mock_client.get_info.side_effect = blocking_get_info
    mock_client.fw_update.side_effect = lamp_flashes

    task1 = hass.async_create_task(install(hass), name="install #1")
    async with asyncio.timeout(5):
        await entered.wait()

    # Install #1 is blocked inside its pre-flag get_info() re-check: the lock is held, but
    # in_progress/ota_in_progress are still untouched.
    assert init_integration.runtime_data.ota_in_progress is False
    assert hass.states.get(entity_id).attributes[ATTR_IN_PROGRESS] is False

    with pytest.raises(HomeAssistantError) as raised:
        await install(hass)
    assert raised.value.translation_key == "update_busy"
    assert init_integration.runtime_data.ota_in_progress is False
    assert hass.states.get(entity_id).attributes[ATTR_IN_PROGRESS] is False

    release.set()
    async with asyncio.timeout(5):
        await task1
    assert init_integration.runtime_data.ota_in_progress is False
    assert hass.states.get(entity_id).attributes[ATTR_IN_PROGRESS] is False
    assert hass.states.get(entity_id).attributes[ATTR_INSTALLED_VERSION] == NEWER


async def test_removing_the_entry_clears_the_wifi_issue(
    hass: HomeAssistant, init_integration, firmware_dir: Path
) -> None:
    put(firmware_dir, "nocreds.img", version=NEWEST, content=CONTENT_WITHOUT_SSID)
    await rescan(hass)
    assert wifi_issue(hass, init_integration) is not None
    assert await hass.config_entries.async_remove(init_integration.entry_id)
    await hass.async_block_till_done()
    assert wifi_issue(hass, init_integration) is None