Add a switch entity for add-ons (#151431)

Co-authored-by: Stefan Agner <stefan@agner.ch>

homeassistant/components/hassio/__init__.py

@@ -73,6 +73,7 @@ from . import (  # noqa: F401
     config_flow,
     diagnostics,
     sensor,
+    switch,
     system_health,
     update,
 )
@@ -149,7 +150,7 @@ _DEPRECATED_HassioServiceInfo = DeprecatedConstant(
 # If new platforms are added, be sure to import them above
 # so we do not make other components that depend on hassio
 # wait for the import of the platforms
-PLATFORMS = [Platform.BINARY_SENSOR, Platform.SENSOR, Platform.UPDATE]
+PLATFORMS = [Platform.BINARY_SENSOR, Platform.SENSOR, Platform.SWITCH, Platform.UPDATE]
 
 CONF_FRONTEND_REPO = "development_repo"
 

homeassistant/components/hassio/coordinator.py

@@ -4,6 +4,7 @@ from __future__ import annotations
 
 import asyncio
 from collections import defaultdict
+from copy import deepcopy
 import logging
 from typing import TYPE_CHECKING, Any
 
@@ -545,3 +546,15 @@ class HassioDataUpdateCoordinator(DataUpdateCoordinator):
         await super()._async_refresh(
             log_failures, raise_on_auth_failed, scheduled, raise_on_entry_error
         )
+
+    async def force_addon_info_data_refresh(self, addon_slug: str) -> None:
+        """Force refresh of addon info data for a specific addon."""
+        try:
+            slug, info = await self._update_addon_info(addon_slug)
+            if info is not None and DATA_KEY_ADDONS in self.data:
+                if slug in self.data[DATA_KEY_ADDONS]:
+                    data = deepcopy(self.data)
+                    data[DATA_KEY_ADDONS][slug].update(info)
+                    self.async_set_updated_data(data)
+        except SupervisorError as err:
+            _LOGGER.warning("Could not refresh info for %s: %s", addon_slug, err)

homeassistant/components/hassio/switch.py

@@ -0,0 +1,90 @@
+"""Switch platform for Hass.io addons."""
+
+from __future__ import annotations
+
+import logging
+from typing import Any
+
+from aiohasupervisor import SupervisorError
+
+from homeassistant.components.switch import SwitchEntity, SwitchEntityDescription
+from homeassistant.config_entries import ConfigEntry
+from homeassistant.const import ATTR_ICON
+from homeassistant.core import HomeAssistant
+from homeassistant.exceptions import HomeAssistantError
+from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
+
+from .const import ADDONS_COORDINATOR, ATTR_STARTED, ATTR_STATE, DATA_KEY_ADDONS
+from .entity import HassioAddonEntity
+from .handler import get_supervisor_client
+
+_LOGGER = logging.getLogger(__name__)
+
+
+ENTITY_DESCRIPTION = SwitchEntityDescription(
+    key=ATTR_STATE,
+    name=None,
+    icon="mdi:puzzle",
+    entity_registry_enabled_default=False,
+)
+
+
+async def async_setup_entry(
+    hass: HomeAssistant,
+    config_entry: ConfigEntry,
+    async_add_entities: AddConfigEntryEntitiesCallback,
+) -> None:
+    """Switch set up for Hass.io config entry."""
+    coordinator = hass.data[ADDONS_COORDINATOR]
+
+    async_add_entities(
+        HassioAddonSwitch(
+            addon=addon,
+            coordinator=coordinator,
+            entity_description=ENTITY_DESCRIPTION,
+        )
+        for addon in coordinator.data[DATA_KEY_ADDONS].values()
+    )
+
+
+class HassioAddonSwitch(HassioAddonEntity, SwitchEntity):
+    """Switch for Hass.io add-ons."""
+
+    @property
+    def is_on(self) -> bool | None:
+        """Return true if the add-on is on."""
+        addon_data = self.coordinator.data[DATA_KEY_ADDONS].get(self._addon_slug, {})
+        state = addon_data.get(self.entity_description.key)
+        return state == ATTR_STARTED
+
+    @property
+    def entity_picture(self) -> str | None:
+        """Return the icon of the add-on if any."""
+        if not self.available:
+            return None
+        addon_data = self.coordinator.data[DATA_KEY_ADDONS].get(self._addon_slug, {})
+        if addon_data.get(ATTR_ICON):
+            return f"/api/hassio/addons/{self._addon_slug}/icon"
+        return None
+
+    async def async_turn_on(self, **kwargs: Any) -> None:
+        """Turn the entity on."""
+        supervisor_client = get_supervisor_client(self.hass)
+        try:
+            await supervisor_client.addons.start_addon(self._addon_slug)
+        except SupervisorError as err:
+            _LOGGER.error("Failed to start addon %s: %s", self._addon_slug, err)
+            raise HomeAssistantError(err) from err
+
+        await self.coordinator.force_addon_info_data_refresh(self._addon_slug)
+
+    async def async_turn_off(self, **kwargs: Any) -> None:
+        """Turn the entity off."""
+        supervisor_client = get_supervisor_client(self.hass)
+        try:
+            await supervisor_client.addons.stop_addon(self._addon_slug)
+        except SupervisorError as err:
+            _LOGGER.error("Failed to stop addon %s: %s", self._addon_slug, err)
+            raise HomeAssistantError(err) from err
+
+        await self.coordinator.force_addon_info_data_refresh(self._addon_slug)

tests/components/hassio/test_switch.py

@@ -0,0 +1,320 @@
+"""The tests for the hassio switch."""
+
+import os
+from unittest.mock import AsyncMock, patch
+
+import pytest
+
+from homeassistant.components.hassio import DOMAIN
+from homeassistant.core import HomeAssistant
+from homeassistant.helpers import entity_registry as er
+from homeassistant.setup import async_setup_component
+
+from .common import MOCK_REPOSITORIES, MOCK_STORE_ADDONS
+
+from tests.common import MockConfigEntry
+from tests.test_util.aiohttp import AiohttpClientMocker
+
+MOCK_ENVIRON = {"SUPERVISOR": "127.0.0.1", "SUPERVISOR_TOKEN": "abcdefgh"}
+
+
+@pytest.fixture(autouse=True)
+def mock_all(
+    aioclient_mock: AiohttpClientMocker,
+    addon_installed: AsyncMock,
+    store_info: AsyncMock,
+    addon_changelog: AsyncMock,
+    addon_stats: AsyncMock,
+    resolution_info: AsyncMock,
+) -> None:
+    """Mock all setup requests."""
+    aioclient_mock.post("http://127.0.0.1/homeassistant/options", json={"result": "ok"})
+    aioclient_mock.post("http://127.0.0.1/supervisor/options", json={"result": "ok"})
+    aioclient_mock.get(
+        "http://127.0.0.1/info",
+        json={
+            "result": "ok",
+            "data": {
+                "supervisor": "222",
+                "homeassistant": "0.110.0",
+                "hassos": "1.2.3",
+            },
+        },
+    )
+    aioclient_mock.get(
+        "http://127.0.0.1/host/info",
+        json={
+            "result": "ok",
+            "data": {
+                "result": "ok",
+                "data": {
+                    "chassis": "vm",
+                    "operating_system": "Debian GNU/Linux 10 (buster)",
+                    "kernel": "4.19.0-6-amd64",
+                },
+            },
+        },
+    )
+    aioclient_mock.get(
+        "http://127.0.0.1/core/info",
+        json={"result": "ok", "data": {"version_latest": "1.0.0", "version": "1.0.0"}},
+    )
+    aioclient_mock.get(
+        "http://127.0.0.1/os/info",
+        json={
+            "result": "ok",
+            "data": {
+                "version_latest": "1.0.0",
+                "version": "1.0.0",
+                "update_available": False,
+            },
+        },
+    )
+    aioclient_mock.get(
+        "http://127.0.0.1/supervisor/info",
+        json={
+            "result": "ok",
+            "data": {
+                "result": "ok",
+                "version": "1.0.0",
+                "version_latest": "1.0.0",
+                "auto_update": True,
+                "addons": [
+                    {
+                        "name": "test",
+                        "state": "started",
+                        "slug": "test",
+                        "installed": True,
+                        "update_available": True,
+                        "icon": False,
+                        "version": "2.0.0",
+                        "version_latest": "2.0.1",
+                        "repository": "core",
+                        "url": "https://github.com/home-assistant/addons/test",
+                    },
+                    {
+                        "name": "test-two",
+                        "state": "stopped",
+                        "slug": "test-two",
+                        "installed": True,
+                        "update_available": False,
+                        "icon": True,
+                        "version": "3.1.0",
+                        "version_latest": "3.1.0",
+                        "repository": "core",
+                        "url": "https://github.com",
+                    },
+                ],
+            },
+        },
+    )
+    aioclient_mock.get(
+        "http://127.0.0.1/core/stats",
+        json={
+            "result": "ok",
+            "data": {
+                "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,
+            },
+        },
+    )
+    aioclient_mock.get(
+        "http://127.0.0.1/supervisor/stats",
+        json={
+            "result": "ok",
+            "data": {
+                "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,
+            },
+        },
+    )
+    aioclient_mock.get(
+        "http://127.0.0.1/ingress/panels", json={"result": "ok", "data": {"panels": {}}}
+    )
+    aioclient_mock.get(
+        "http://127.0.0.1/network/info",
+        json={
+            "result": "ok",
+            "data": {
+                "host_internet": True,
+                "supervisor_internet": True,
+            },
+        },
+    )
+
+
+@pytest.mark.parametrize(
+    ("store_addons", "store_repositories"), [(MOCK_STORE_ADDONS, MOCK_REPOSITORIES)]
+)
+@pytest.mark.parametrize(
+    ("entity_id", "expected", "addon_state"),
+    [
+        ("switch.test", "on", "started"),
+        ("switch.test_two", "off", "stopped"),
+    ],
+)
+async def test_switch_state(
+    hass: HomeAssistant,
+    entity_id: str,
+    expected: str,
+    addon_state: str,
+    aioclient_mock: AiohttpClientMocker,
+    entity_registry: er.EntityRegistry,
+    addon_installed: AsyncMock,
+) -> None:
+    """Test hassio addon switch state."""
+    addon_installed.return_value.state = addon_state
+    config_entry = MockConfigEntry(domain=DOMAIN, data={}, unique_id=DOMAIN)
+    config_entry.add_to_hass(hass)
+
+    with patch.dict(os.environ, MOCK_ENVIRON):
+        result = await async_setup_component(
+            hass,
+            "hassio",
+            {"http": {"server_port": 9999, "server_host": "127.0.0.1"}, "hassio": {}},
+        )
+        assert result
+    await hass.async_block_till_done()
+
+    # Verify that the entity is disabled by default.
+    assert hass.states.get(entity_id) is None
+
+    # Enable the entity.
+    entity_registry.async_update_entity(entity_id, disabled_by=None)
+    await hass.config_entries.async_reload(config_entry.entry_id)
+    await hass.async_block_till_done()
+
+    # Verify that the entity have the expected state.
+    state = hass.states.get(entity_id)
+    assert state is not None
+    assert state.state == expected
+
+
+@pytest.mark.parametrize(
+    ("store_addons", "store_repositories"), [(MOCK_STORE_ADDONS, MOCK_REPOSITORIES)]
+)
+async def test_switch_turn_on(
+    hass: HomeAssistant,
+    aioclient_mock: AiohttpClientMocker,
+    entity_registry: er.EntityRegistry,
+    addon_installed: AsyncMock,
+) -> None:
+    """Test turning on addon switch."""
+    entity_id = "switch.test_two"
+    addon_installed.return_value.state = "stopped"
+
+    # Mock the start addon API call
+    aioclient_mock.post("http://127.0.0.1/addons/test-two/start", json={"result": "ok"})
+
+    config_entry = MockConfigEntry(domain=DOMAIN, data={}, unique_id=DOMAIN)
+    config_entry.add_to_hass(hass)
+
+    with patch.dict(os.environ, MOCK_ENVIRON):
+        result = await async_setup_component(
+            hass,
+            "hassio",
+            {"http": {"server_port": 9999, "server_host": "127.0.0.1"}, "hassio": {}},
+        )
+        assert result
+    await hass.async_block_till_done()
+
+    # Verify that the entity is disabled by default.
+    assert hass.states.get(entity_id) is None
+
+    # Enable the entity.
+    entity_registry.async_update_entity(entity_id, disabled_by=None)
+    await hass.config_entries.async_reload(config_entry.entry_id)
+    await hass.async_block_till_done()
+
+    # Verify initial state is off
+    state = hass.states.get(entity_id)
+    assert state is not None
+    assert state.state == "off"
+
+    # Turn on the switch
+    await hass.services.async_call(
+        "switch",
+        "turn_on",
+        {"entity_id": entity_id},
+        blocking=True,
+    )
+
+    # Verify the API was called
+    assert len(aioclient_mock.mock_calls) > 0
+    start_call_found = False
+    for call in aioclient_mock.mock_calls:
+        if call[1].path == "/addons/test-two/start" and call[0] == "POST":
+            start_call_found = True
+            break
+    assert start_call_found
+
+
+@pytest.mark.parametrize(
+    ("store_addons", "store_repositories"), [(MOCK_STORE_ADDONS, MOCK_REPOSITORIES)]
+)
+async def test_switch_turn_off(
+    hass: HomeAssistant,
+    aioclient_mock: AiohttpClientMocker,
+    entity_registry: er.EntityRegistry,
+    addon_installed: AsyncMock,
+) -> None:
+    """Test turning off addon switch."""
+    entity_id = "switch.test"
+    addon_installed.return_value.state = "started"
+
+    # Mock the stop addon API call
+    aioclient_mock.post("http://127.0.0.1/addons/test/stop", json={"result": "ok"})
+
+    config_entry = MockConfigEntry(domain=DOMAIN, data={}, unique_id=DOMAIN)
+    config_entry.add_to_hass(hass)
+
+    with patch.dict(os.environ, MOCK_ENVIRON):
+        result = await async_setup_component(
+            hass,
+            "hassio",
+            {"http": {"server_port": 9999, "server_host": "127.0.0.1"}, "hassio": {}},
+        )
+        assert result
+    await hass.async_block_till_done()
+
+    # Verify that the entity is disabled by default.
+    assert hass.states.get(entity_id) is None
+
+    # Enable the entity.
+    entity_registry.async_update_entity(entity_id, disabled_by=None)
+    await hass.config_entries.async_reload(config_entry.entry_id)
+    await hass.async_block_till_done()
+
+    # Verify initial state is on
+    state = hass.states.get(entity_id)
+    assert state is not None
+    assert state.state == "on"
+
+    # Turn off the switch
+    await hass.services.async_call(
+        "switch",
+        "turn_off",
+        {"entity_id": entity_id},
+        blocking=True,
+    )
+
+    # Verify the API was called
+    assert len(aioclient_mock.mock_calls) > 0
+    stop_call_found = False
+    for call in aioclient_mock.mock_calls:
+        if call[1].path == "/addons/test/stop" and call[0] == "POST":
+            stop_call_found = True
+            break
+    assert stop_call_found