Refactor diagnostics, create backup and green/yellow settings from handler (home-assistant#154098)

Co-authored-by: Stefan Agner <[email protected]>

Author: mdegat01

homeassistant/components/hassio/__init__.py

@@ -13,6 +13,7 @@
 from typing import Any, NamedTuple
 
 from aiohasupervisor import SupervisorError
+from aiohasupervisor.models import GreenOptions, YellowOptions  # noqa: F401
 import voluptuous as vol
 
 from homeassistant.auth.const import GROUP_ID_ADMIN
@@ -123,11 +124,6 @@
 from .handler import (  # noqa: F401
     HassIO,
     HassioAPIError,
-    async_create_backup,
-    async_get_green_settings,
-    async_get_yellow_settings,
-    async_set_green_settings,
-    async_set_yellow_settings,
     async_update_diagnostics,
     get_supervisor_client,
 )

homeassistant/components/hassio/addon_manager.py

@@ -15,13 +15,14 @@
     AddonsOptions,
     AddonState as SupervisorAddonState,
     InstalledAddonComplete,
+    PartialBackupOptions,
     StoreAddonUpdate,
 )
 
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.exceptions import HomeAssistantError
 
-from .handler import HassioAPIError, async_create_backup, get_supervisor_client
+from .handler import HassioAPIError, get_supervisor_client
 
 type _FuncType[_T, **_P, _R] = Callable[Concatenate[_T, _P], Awaitable[_R]]
 type _ReturnFuncType[_T, **_P, _R] = Callable[
@@ -261,17 +262,18 @@ async def async_stop_addon(self) -> None:
         """Stop the managed add-on."""
         await self._supervisor_client.addons.stop_addon(self.addon_slug)
 
-    @api_error("Failed to create a backup of the {addon_name} add-on")
+    @api_error(
+        "Failed to create a backup of the {addon_name} add-on",
+        expected_error_type=SupervisorError,
+    )
     async def async_create_backup(self) -> None:
         """Create a partial backup of the managed add-on."""
         addon_info = await self.async_get_addon_info()
         name = f"addon_{self.addon_slug}_{addon_info.version}"
 
         self._logger.debug("Creating backup: %s", name)
-        await async_create_backup(
-            self._hass,
-            {"name": name, "addons": [self.addon_slug]},
-            partial=True,
+        await self._supervisor_client.backups.partial_backup(
+            PartialBackupOptions(name=name, addons={self.addon_slug})
         )
 
     async def async_configure_addon(

homeassistant/components/hassio/handler.py

@@ -10,6 +10,7 @@
 from typing import Any
 
 from aiohasupervisor import SupervisorClient
+from aiohasupervisor.models import SupervisorOptions
 import aiohttp
 from yarl import URL
 
@@ -22,7 +23,6 @@
 from homeassistant.const import SERVER_PORT
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.singleton import singleton
-from homeassistant.loader import bind_hass
 
 from .const import ATTR_MESSAGE, ATTR_RESULT, DATA_COMPONENT, X_HASS_SOURCE
 
@@ -66,73 +66,6 @@ async def _wrapper(*argv: _P.args, **kwargs: _P.kwargs) -> Any:
     return _wrapper
 
 
-@bind_hass
-async def async_update_diagnostics(hass: HomeAssistant, diagnostics: bool) -> bool:
-    """Update Supervisor diagnostics toggle.
-
-    The caller of the function should handle HassioAPIError.
-    """
-    hassio = hass.data[DATA_COMPONENT]
-    return await hassio.update_diagnostics(diagnostics)
-
-
-@bind_hass
-@api_data
-async def async_create_backup(
-    hass: HomeAssistant, payload: dict, partial: bool = False
-) -> dict:
-    """Create a full or partial backup.
-
-    The caller of the function should handle HassioAPIError.
-    """
-    hassio = hass.data[DATA_COMPONENT]
-    backup_type = "partial" if partial else "full"
-    command = f"/backups/new/{backup_type}"
-    return await hassio.send_command(command, payload=payload, timeout=None)
-
-
-@api_data
-async def async_get_green_settings(hass: HomeAssistant) -> dict[str, bool]:
-    """Return settings specific to Home Assistant Green."""
-    hassio = hass.data[DATA_COMPONENT]
-    return await hassio.send_command("/os/boards/green", method="get")
-
-
-@api_data
-async def async_set_green_settings(
-    hass: HomeAssistant, settings: dict[str, bool]
-) -> dict:
-    """Set settings specific to Home Assistant Green.
-
-    Returns an empty dict.
-    """
-    hassio = hass.data[DATA_COMPONENT]
-    return await hassio.send_command(
-        "/os/boards/green", method="post", payload=settings
-    )
-
-
-@api_data
-async def async_get_yellow_settings(hass: HomeAssistant) -> dict[str, bool]:
-    """Return settings specific to Home Assistant Yellow."""
-    hassio = hass.data[DATA_COMPONENT]
-    return await hassio.send_command("/os/boards/yellow", method="get")
-
-
-@api_data
-async def async_set_yellow_settings(
-    hass: HomeAssistant, settings: dict[str, bool]
-) -> dict:
-    """Set settings specific to Home Assistant Yellow.
-
-    Returns an empty dict.
-    """
-    hassio = hass.data[DATA_COMPONENT]
-    return await hassio.send_command(
-        "/os/boards/yellow", method="post", payload=settings
-    )
-
-
 class HassIO:
     """Small API wrapper for Hass.io."""
 
@@ -257,16 +190,6 @@ def update_hass_config(self, timezone: str, country: str | None) -> Coroutine:
             "/supervisor/options", payload={"timezone": timezone, "country": country}
         )
 
-    @_api_bool
-    def update_diagnostics(self, diagnostics: bool) -> Coroutine:
-        """Update Supervisor diagnostics setting.
-
-        This method returns a coroutine.
-        """
-        return self.send_command(
-            "/supervisor/options", payload={"diagnostics": diagnostics}
-        )
-
     async def send_command(
         self,
         command: str,
@@ -341,3 +264,13 @@ def get_supervisor_client(hass: HomeAssistant) -> SupervisorClient:
         os.environ.get("SUPERVISOR_TOKEN", ""),
         session=hassio.websession,
     )
+
+
+async def async_update_diagnostics(hass: HomeAssistant, diagnostics: bool) -> None:
+    """Update Supervisor diagnostics toggle.
+
+    The caller of the function should handle SupervisorError.
+    """
+    await get_supervisor_client(hass).supervisor.set_options(
+        SupervisorOptions(diagnostics=diagnostics)
+    )

homeassistant/components/homeassistant_green/config_flow.py

@@ -6,21 +6,20 @@
 import logging
 from typing import Any
 
-import aiohttp
 import voluptuous as vol
 
 from homeassistant.components.hassio import (
-    HassioAPIError,
-    async_get_green_settings,
-    async_set_green_settings,
+    GreenOptions,
+    SupervisorError,
+    get_supervisor_client,
 )
 from homeassistant.config_entries import (
     ConfigEntry,
     ConfigFlow,
     ConfigFlowResult,
     OptionsFlow,
 )
-from homeassistant.core import callback
+from homeassistant.core import HomeAssistant, async_get_hass, callback
 from homeassistant.helpers import selector
 from homeassistant.helpers.hassio import is_hassio
 
@@ -49,7 +48,7 @@ def async_get_options_flow(
         config_entry: ConfigEntry,
     ) -> HomeAssistantGreenOptionsFlow:
         """Return the options flow."""
-        return HomeAssistantGreenOptionsFlow()
+        return HomeAssistantGreenOptionsFlow(async_get_hass())
 
     async def async_step_system(
         self, data: dict[str, Any] | None = None
@@ -63,6 +62,11 @@ class HomeAssistantGreenOptionsFlow(OptionsFlow):
 
     _hw_settings: dict[str, bool] | None = None
 
+    def __init__(self, hass: HomeAssistant, *args: Any, **kwargs: Any) -> None:
+        """Instantiate options flow."""
+        super().__init__(*args, **kwargs)
+        self._supervisor_client = get_supervisor_client(hass)
+
     async def async_step_init(
         self, user_input: dict[str, Any] | None = None
     ) -> ConfigFlowResult:
@@ -76,27 +80,27 @@ async def async_step_hardware_settings(
         self, user_input: dict[str, Any] | None = None
     ) -> ConfigFlowResult:
         """Handle hardware settings."""
-
         if user_input is not None:
             if self._hw_settings == user_input:
                 return self.async_create_entry(data={})
             try:
                 async with asyncio.timeout(10):
-                    await async_set_green_settings(self.hass, user_input)
-            except (aiohttp.ClientError, TimeoutError, HassioAPIError) as err:
+                    await self._supervisor_client.os.set_green_options(
+                        GreenOptions.from_dict(user_input)
+                    )
+            except (TimeoutError, SupervisorError) as err:
                 _LOGGER.warning("Failed to write hardware settings", exc_info=err)
                 return self.async_abort(reason="write_hw_settings_error")
             return self.async_create_entry(data={})
 
         try:
             async with asyncio.timeout(10):
-                self._hw_settings: dict[str, bool] = await async_get_green_settings(
-                    self.hass
-                )
-        except (aiohttp.ClientError, TimeoutError, HassioAPIError) as err:
+                green_info = await self._supervisor_client.os.green_info()
+        except (TimeoutError, SupervisorError) as err:
             _LOGGER.warning("Failed to read hardware settings", exc_info=err)
             return self.async_abort(reason="read_hw_settings_error")
 
+        self._hw_settings: dict[str, bool] = green_info.to_dict()
         schema = self.add_suggested_values_to_schema(
             STEP_HW_SETTINGS_SCHEMA, self._hw_settings
         )

homeassistant/components/homeassistant_yellow/config_flow.py

@@ -7,13 +7,11 @@
 import logging
 from typing import TYPE_CHECKING, Any, Protocol, final
 
-import aiohttp
 import voluptuous as vol
 
 from homeassistant.components.hassio import (
-    HassioAPIError,
-    async_get_yellow_settings,
-    async_set_yellow_settings,
+    SupervisorError,
+    YellowOptions,
     get_supervisor_client,
 )
 from homeassistant.components.homeassistant_hardware.firmware_config_flow import (
@@ -222,21 +220,22 @@ async def async_step_hardware_settings(
                 return self.async_create_entry(data={})
             try:
                 async with asyncio.timeout(10):
-                    await async_set_yellow_settings(self.hass, user_input)
-            except (aiohttp.ClientError, TimeoutError, HassioAPIError) as err:
+                    await self._supervisor_client.os.set_yellow_options(
+                        YellowOptions.from_dict(user_input)
+                    )
+            except (TimeoutError, SupervisorError) as err:
                 _LOGGER.warning("Failed to write hardware settings", exc_info=err)
                 return self.async_abort(reason="write_hw_settings_error")
             return await self.async_step_reboot_menu()
 
         try:
             async with asyncio.timeout(10):
-                self._hw_settings: dict[str, bool] = await async_get_yellow_settings(
-                    self.hass
-                )
-        except (aiohttp.ClientError, TimeoutError, HassioAPIError) as err:
+                yellow_info = await self._supervisor_client.os.yellow_info()
+        except (TimeoutError, SupervisorError) as err:
             _LOGGER.warning("Failed to read hardware settings", exc_info=err)
             return self.async_abort(reason="read_hw_settings_error")
 
+        self._hw_settings: dict[str, bool] = yellow_info.to_dict()
         schema = self.add_suggested_values_to_schema(
             STEP_HW_SETTINGS_SCHEMA, self._hw_settings
         )

tests/components/conftest.py

@@ -14,11 +14,13 @@
 
 from aiohasupervisor.models import (
     Discovery,
+    GreenInfo,
     JobsInfo,
     Repository,
     ResolutionInfo,
     StoreAddon,
     StoreInfo,
+    YellowInfo,
 )
 import pytest
 import voluptuous as vol
@@ -430,11 +432,9 @@ def uninstall_addon_fixture(supervisor_client: AsyncMock) -> AsyncMock:
 
 
 @pytest.fixture(name="create_backup")
-def create_backup_fixture() -> Generator[AsyncMock]:
+def create_backup_fixture(supervisor_client: AsyncMock) -> AsyncMock:
     """Mock create backup."""
-    from .hassio.common import mock_create_backup  # noqa: PLC0415
-
-    yield from mock_create_backup()
+    return supervisor_client.backups.partial_backup
 
 
 @pytest.fixture(name="update_addon")
@@ -517,6 +517,24 @@ def jobs_info_fixture(supervisor_client: AsyncMock) -> AsyncMock:
     return supervisor_client.jobs.info
 
 
+@pytest.fixture(name="os_yellow_info")
+def os_yellow_info_fixture(supervisor_client: AsyncMock) -> AsyncMock:
+    """Mock yellow info API from supervisor OS."""
+    supervisor_client.os.yellow_info.return_value = YellowInfo(
+        disk_led=True, heartbeat_led=True, power_led=True
+    )
+    return supervisor_client.os.yellow_info
+
+
+@pytest.fixture(name="os_green_info")
+def os_green_info_fixture(supervisor_client: AsyncMock) -> AsyncMock:
+    """Mock green info API from supervisor OS."""
+    supervisor_client.os.green_info.return_value = GreenInfo(
+        activity_led=True, power_led=True, system_health_led=True
+    )
+    return supervisor_client.os.green_info
+
+
 @pytest.fixture(name="supervisor_client")
 def supervisor_client() -> Generator[AsyncMock]:
     """Mock the supervisor client."""

tests/components/hassio/common.py

@@ -2,12 +2,11 @@
 
 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 unittest.mock import AsyncMock, Mock
 
 from aiohasupervisor.models import (
     AddonsOptions,
@@ -197,14 +196,6 @@ async def set_addon_options(slug: str, options: AddonsOptions) -> None:
     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(