core/tests/components/hassio/common.py

"""Provide common test tools for hassio."""

from __future__ import annotations

from collections.abc import Generator
from dataclasses import fields
import logging
from types import MethodType
from typing import Any
from unittest.mock import AsyncMock, Mock, patch

from aiohasupervisor.models import (
    AddonsOptions,
    AddonsStats,
    AddonStage,
    InstalledAddonComplete,
    Repository,
    StoreAddon,
    StoreAddonComplete,
)

from homeassistant.components.hassio.addon_manager import AddonManager
from homeassistant.core import HomeAssistant

LOGGER = logging.getLogger(__name__)
INSTALLED_ADDON_FIELDS = [field.name for field in fields(InstalledAddonComplete)]
STORE_ADDON_FIELDS = [field.name for field in fields(StoreAddonComplete)]
ADDONS_STATS_FIELDS = [field.name for field in fields(AddonsStats)]

MOCK_STORE_ADDONS = [
    StoreAddon(
        name="test",
        arch=[],
        documentation=False,
        advanced=False,
        available=True,
        build=False,
        description="Test add-on service",
        homeassistant=None,
        icon=False,
        logo=False,
        repository="core",
        slug="core_test",
        stage=AddonStage.EXPERIMENTAL,
        update_available=False,
        url="https://example.com/addons/tree/master/test",
        version_latest="1.0.0",
        version="1.0.0",
        installed=True,
    )
]

MOCK_REPOSITORIES = [
    Repository(
        slug="core",
        name="Official add-ons",
        source="core",
        url="https://home-assistant.io/addons",
        maintainer="Home Assistant",
    )
]


def mock_to_dict(obj: Mock, fields: list[str]) -> dict[str, Any]:
    """Aiohasupervisor mocks to dictionary representation."""
    return {
        field: getattr(obj, field)
        for field in fields
        if not isinstance(getattr(obj, field), Mock)
    }


def mock_addon_manager(hass: HomeAssistant) -> AddonManager:
    """Return an AddonManager instance."""
    return AddonManager(hass, LOGGER, "Test", "test_addon")


def mock_addon_store_info(
    supervisor_client: AsyncMock,
    addon_store_info_side_effect: Any | None,
) -> AsyncMock:
    """Mock Supervisor add-on store info."""
    supervisor_client.store.addon_info.side_effect = addon_store_info_side_effect

    supervisor_client.store.addon_info.return_value = addon_info = Mock(
        spec=StoreAddonComplete,
        slug="test",
        repository="core",
        available=True,
        installed=False,
        update_available=False,
        version="1.0.0",
        supervisor_api=False,
        supervisor_role="default",
    )
    addon_info.name = "test"
    addon_info.to_dict = MethodType(
        lambda self: mock_to_dict(self, STORE_ADDON_FIELDS),
        addon_info,
    )
    return supervisor_client.store.addon_info


def mock_addon_info(
    supervisor_client: AsyncMock, addon_info_side_effect: Any | None
) -> AsyncMock:
    """Mock Supervisor add-on info."""
    supervisor_client.addons.addon_info.side_effect = addon_info_side_effect

    supervisor_client.addons.addon_info.return_value = addon_info = Mock(
        spec=InstalledAddonComplete,
        slug="test",
        repository="core",
        available=False,
        hostname="",
        options={},
        state="unknown",
        update_available=False,
        version=None,
        supervisor_api=False,
        supervisor_role="default",
    )
    addon_info.name = "test"
    addon_info.to_dict = MethodType(
        lambda self: mock_to_dict(self, INSTALLED_ADDON_FIELDS),
        addon_info,
    )
    return supervisor_client.addons.addon_info


def mock_addon_not_installed(
    addon_store_info: AsyncMock, addon_info: AsyncMock
) -> AsyncMock:
    """Mock add-on not installed."""
    addon_store_info.return_value.available = True
    return addon_info


def mock_addon_installed(
    addon_store_info: AsyncMock, addon_info: AsyncMock
) -> AsyncMock:
    """Mock add-on already installed but not running."""
    addon_store_info.return_value.available = True
    addon_store_info.return_value.installed = True
    addon_info.return_value.available = True
    addon_info.return_value.hostname = "core-test-addon"
    addon_info.return_value.state = "stopped"
    addon_info.return_value.version = "1.0.0"
    return addon_info


def mock_addon_running(addon_store_info: AsyncMock, addon_info: AsyncMock) -> AsyncMock:
    """Mock add-on already running."""
    addon_store_info.return_value.available = True
    addon_store_info.return_value.installed = True
    addon_info.return_value.state = "started"
    return addon_info


def mock_install_addon_side_effect(
    addon_store_info: AsyncMock, addon_info: AsyncMock
) -> Any | None:
    """Return the install add-on side effect."""

    async def install_addon(addon: str):
        """Mock install add-on."""
        addon_store_info.return_value.available = True
        addon_store_info.return_value.installed = True
        addon_info.return_value.available = True
        addon_info.return_value.state = "stopped"
        addon_info.return_value.version = "1.0.0"

    return install_addon


def mock_start_addon_side_effect(
    addon_store_info: AsyncMock, addon_info: AsyncMock
) -> Any | None:
    """Return the start add-on options side effect."""

    async def start_addon(addon: str) -> None:
        """Mock start add-on."""
        addon_store_info.return_value.available = True
        addon_store_info.return_value.installed = True
        addon_info.return_value.available = True
        addon_info.return_value.state = "started"

    return start_addon


def mock_set_addon_options_side_effect(addon_options: dict[str, Any]) -> Any | None:
    """Return the set add-on options side effect."""

    async def set_addon_options(slug: str, options: AddonsOptions) -> None:
        """Mock set add-on options."""
        addon_options.update(options.config)

    return set_addon_options


def mock_create_backup() -> Generator[AsyncMock]:
    """Mock create backup."""
    with patch(
        "homeassistant.components.hassio.addon_manager.async_create_backup"
    ) as create_backup:
        yield create_backup


def mock_addon_stats(supervisor_client: AsyncMock) -> AsyncMock:
    """Mock addon stats."""
    supervisor_client.addons.addon_stats.return_value = addon_stats = Mock(
        spec=AddonsStats,
        cpu_percent=0.99,
        memory_usage=182611968,
        memory_limit=3977146368,
        memory_percent=4.59,
        network_rx=362570232,
        network_tx=82374138,
        blk_read=46010945536,
        blk_write=15051526144,
    )
    addon_stats.to_dict = MethodType(
        lambda self: mock_to_dict(self, ADDONS_STATS_FIELDS),
        addon_stats,
    )
    return supervisor_client.addons.addon_stats