core/tests/components/hassio/common.py

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

from __future__ import annotations

from collections.abc import Generator
import logging
from typing import Any
from unittest.mock import DEFAULT, AsyncMock, patch

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

LOGGER = logging.getLogger(__name__)


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


def mock_discovery_info() -> Any:
    """Return the discovery info from the supervisor."""
    return DEFAULT


def mock_get_addon_discovery_info(
    discovery_info: dict[str, Any], discovery_info_side_effect: Any | None
) -> Generator[AsyncMock]:
    """Mock get add-on discovery info."""
    with patch(
        "homeassistant.components.hassio.addon_manager.async_get_addon_discovery_info",
        side_effect=discovery_info_side_effect,
        return_value=discovery_info,
    ) as get_addon_discovery_info:
        yield get_addon_discovery_info


def mock_addon_store_info(
    addon_store_info_side_effect: Any | None,
) -> Generator[AsyncMock]:
    """Mock Supervisor add-on store info."""
    with patch(
        "homeassistant.components.hassio.addon_manager.async_get_addon_store_info",
        side_effect=addon_store_info_side_effect,
    ) as addon_store_info:
        addon_store_info.return_value = {
            "available": True,
            "installed": None,
            "state": None,
            "version": "1.0.0",
        }
        yield addon_store_info


def mock_addon_info(addon_info_side_effect: Any | None) -> Generator[AsyncMock]:
    """Mock Supervisor add-on info."""
    with patch(
        "homeassistant.components.hassio.addon_manager.async_get_addon_info",
        side_effect=addon_info_side_effect,
    ) as addon_info:
        addon_info.return_value = {
            "available": False,
            "hostname": None,
            "options": {},
            "state": None,
            "update_available": False,
            "version": None,
        }
        yield 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,
        "installed": "1.0.0",
        "state": "stopped",
        "version": "1.0.0",
    }
    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,
        "installed": "1.0.0",
        "state": "started",
        "version": "1.0.0",
    }
    addon_info.return_value["available"] = True
    addon_info.return_value["hostname"] = "core-test-addon"
    addon_info.return_value["state"] = "started"
    addon_info.return_value["version"] = "1.0.0"
    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(hass: HomeAssistant, slug):
        """Mock install add-on."""
        addon_store_info.return_value = {
            "available": True,
            "installed": "1.0.0",
            "state": "stopped",
            "version": "1.0.0",
        }
        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_install_addon(install_addon_side_effect: Any | None) -> Generator[AsyncMock]:
    """Mock install add-on."""

    with patch(
        "homeassistant.components.hassio.addon_manager.async_install_addon",
        side_effect=install_addon_side_effect,
    ) as install_addon:
        yield 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(hass: HomeAssistant, slug):
        """Mock start add-on."""
        addon_store_info.return_value = {
            "available": True,
            "installed": "1.0.0",
            "state": "started",
            "version": "1.0.0",
        }
        addon_info.return_value["available"] = True
        addon_info.return_value["state"] = "started"

    return start_addon


def mock_start_addon(start_addon_side_effect: Any | None) -> Generator[AsyncMock]:
    """Mock start add-on."""
    with patch(
        "homeassistant.components.hassio.addon_manager.async_start_addon",
        side_effect=start_addon_side_effect,
    ) as start_addon:
        yield start_addon


def mock_stop_addon() -> Generator[AsyncMock]:
    """Mock stop add-on."""
    with patch(
        "homeassistant.components.hassio.addon_manager.async_stop_addon"
    ) as stop_addon:
        yield stop_addon


def mock_restart_addon(restart_addon_side_effect: Any | None) -> Generator[AsyncMock]:
    """Mock restart add-on."""
    with patch(
        "homeassistant.components.hassio.addon_manager.async_restart_addon",
        side_effect=restart_addon_side_effect,
    ) as restart_addon:
        yield restart_addon


def mock_uninstall_addon() -> Generator[AsyncMock]:
    """Mock uninstall add-on."""
    with patch(
        "homeassistant.components.hassio.addon_manager.async_uninstall_addon"
    ) as uninstall_addon:
        yield uninstall_addon


def mock_addon_options(addon_info: AsyncMock) -> dict[str, Any]:
    """Mock add-on options."""
    return addon_info.return_value["options"]


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(hass: HomeAssistant, slug: str, options: dict) -> None:
        """Mock set add-on options."""
        addon_options.update(options["options"])

    return set_addon_options


def mock_set_addon_options(
    set_addon_options_side_effect: Any | None,
) -> Generator[AsyncMock]:
    """Mock set add-on options."""
    with patch(
        "homeassistant.components.hassio.addon_manager.async_set_addon_options",
        side_effect=set_addon_options_side_effect,
    ) as set_options:
        yield set_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_update_addon() -> Generator[AsyncMock]:
    """Mock update add-on."""
    with patch(
        "homeassistant.components.hassio.addon_manager.async_update_addon"
    ) as update_addon:
        yield update_addon