Add backup support to the hassio OS update entity (#142580)

* Add backup support to the hassio OS update entity

* Remove meaningless assert

homeassistant/components/hassio/update.py

@@ -5,7 +5,6 @@ from __future__ import annotations
 from typing import Any
 
 from aiohasupervisor import SupervisorError
-from aiohasupervisor.models import OSUpdate
 from awesomeversion import AwesomeVersion, AwesomeVersionStrategy
 
 from homeassistant.components.update import (
@@ -36,7 +35,7 @@ from .entity import (
     HassioOSEntity,
     HassioSupervisorEntity,
 )
-from .update_helper import update_addon, update_core
+from .update_helper import update_addon, update_core, update_os
 
 ENTITY_DESCRIPTION = UpdateEntityDescription(
     translation_key="update",
@@ -170,7 +169,9 @@ class SupervisorOSUpdateEntity(HassioOSEntity, UpdateEntity):
     """Update entity to handle updates for the Home Assistant Operating System."""
 
     _attr_supported_features = (
-        UpdateEntityFeature.INSTALL | UpdateEntityFeature.SPECIFIC_VERSION
+        UpdateEntityFeature.INSTALL
+        | UpdateEntityFeature.SPECIFIC_VERSION
+        | UpdateEntityFeature.BACKUP
     )
     _attr_title = "Home Assistant Operating System"
 
@@ -203,14 +204,7 @@ class SupervisorOSUpdateEntity(HassioOSEntity, UpdateEntity):
         self, version: str | None, backup: bool, **kwargs: Any
     ) -> None:
         """Install an update."""
-        try:
+        await update_os(self.hass, version, backup)
-            await self.coordinator.supervisor_client.os.update(
-                OSUpdate(version=version)
-            )
-        except SupervisorError as err:
-            raise HomeAssistantError(
-                f"Error updating Home Assistant Operating System: {err}"
-            ) from err
 
 
 class SupervisorSupervisorUpdateEntity(HassioSupervisorEntity, UpdateEntity):

homeassistant/components/hassio/update_helper.py

@@ -3,7 +3,11 @@
 from __future__ import annotations
 
 from aiohasupervisor import SupervisorError
-from aiohasupervisor.models import HomeAssistantUpdateOptions, StoreAddonUpdate
+from aiohasupervisor.models import (
+    HomeAssistantUpdateOptions,
+    OSUpdate,
+    StoreAddonUpdate,
+)
 
 from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import HomeAssistantError
@@ -57,3 +61,24 @@ async def update_core(hass: HomeAssistant, version: str | None, backup: bool) ->
         )
     except SupervisorError as err:
         raise HomeAssistantError(f"Error updating Home Assistant Core: {err}") from err
+
+
+async def update_os(hass: HomeAssistant, version: str | None, backup: bool) -> None:
+    """Update OS.
+
+    Optionally make a core backup before updating.
+    """
+    client = get_supervisor_client(hass)
+
+    if backup:
+        # pylint: disable-next=import-outside-toplevel
+        from .backup import backup_core_before_update
+
+        await backup_core_before_update(hass)
+
+    try:
+        await client.os.update(OSUpdate(version=version))
+    except SupervisorError as err:
+        raise HomeAssistantError(
+            f"Error updating Home Assistant Operating System: {err}"
+        ) from err

tests/components/hassio/test_update.py

@@ -6,7 +6,11 @@ from typing import Any
 from unittest.mock import AsyncMock, MagicMock, patch
 
 from aiohasupervisor import SupervisorBadRequestError, SupervisorError
-from aiohasupervisor.models import HomeAssistantUpdateOptions, StoreAddonUpdate
+from aiohasupervisor.models import (
+    HomeAssistantUpdateOptions,
+    OSUpdate,
+    StoreAddonUpdate,
+)
 import pytest
 
 from homeassistant.components.backup import BackupManagerError, ManagerBackup
@@ -475,13 +479,123 @@ async def test_update_os(hass: HomeAssistant, supervisor_client: AsyncMock) -> N
     await hass.async_block_till_done()
 
     supervisor_client.os.update.return_value = None
-    await hass.services.async_call(
+    with patch(
-        "update",
+        "homeassistant.components.backup.manager.BackupManager.async_create_backup",
-        "install",
+    ) as mock_create_backup:
-        {"entity_id": "update.home_assistant_operating_system_update"},
+        await hass.services.async_call(
-        blocking=True,
+            "update",
-    )
+            "install",
-    supervisor_client.os.update.assert_called_once()
+            {"entity_id": "update.home_assistant_operating_system_update"},
+            blocking=True,
+        )
+    mock_create_backup.assert_not_called()
+    supervisor_client.os.update.assert_called_once_with(OSUpdate(version=None))
+
+
+@pytest.mark.parametrize(
+    ("commands", "default_mount", "expected_kwargs"),
+    [
+        (
+            [],
+            None,
+            {
+                "agent_ids": ["hassio.local"],
+                "include_addons": None,
+                "include_all_addons": False,
+                "include_database": True,
+                "include_folders": None,
+                "include_homeassistant": True,
+                "name": f"Home Assistant Core {HAVERSION}",
+                "password": None,
+            },
+        ),
+        (
+            [],
+            "my_nas",
+            {
+                "agent_ids": ["hassio.my_nas"],
+                "include_addons": None,
+                "include_all_addons": False,
+                "include_database": True,
+                "include_folders": None,
+                "include_homeassistant": True,
+                "name": f"Home Assistant Core {HAVERSION}",
+                "password": None,
+            },
+        ),
+        (
+            [
+                {
+                    "type": "backup/config/update",
+                    "create_backup": {
+                        "agent_ids": ["test-agent"],
+                        "include_addons": ["my-addon"],
+                        "include_all_addons": True,
+                        "include_database": False,
+                        "include_folders": ["share"],
+                        "name": "cool_backup",
+                        "password": "hunter2",
+                    },
+                },
+            ],
+            None,
+            {
+                "agent_ids": ["test-agent"],
+                "include_addons": ["my-addon"],
+                "include_all_addons": True,
+                "include_database": False,
+                "include_folders": ["share"],
+                "include_homeassistant": True,
+                "name": "cool_backup",
+                "password": "hunter2",
+                "with_automatic_settings": True,
+            },
+        ),
+    ],
+)
+async def test_update_os_with_backup(
+    hass: HomeAssistant,
+    hass_ws_client: WebSocketGenerator,
+    supervisor_client: AsyncMock,
+    commands: list[dict[str, Any]],
+    default_mount: str | None,
+    expected_kwargs: dict[str, Any],
+) -> None:
+    """Test updating OS update entity."""
+    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 setup_backup_integration(hass)
+
+    client = await hass_ws_client(hass)
+    for command in commands:
+        await client.send_json_auto_id(command)
+        result = await client.receive_json()
+        assert result["success"]
+
+    supervisor_client.os.update.return_value = None
+    supervisor_client.mounts.info.return_value.default_backup_mount = default_mount
+    with patch(
+        "homeassistant.components.backup.manager.BackupManager.async_create_backup",
+    ) as mock_create_backup:
+        await hass.services.async_call(
+            "update",
+            "install",
+            {
+                "entity_id": "update.home_assistant_operating_system_update",
+                "backup": True,
+            },
+            blocking=True,
+        )
+    mock_create_backup.assert_called_once_with(**expected_kwargs)
+    supervisor_client.os.update.assert_called_once_with(OSUpdate(version=None))
 
 
 async def test_update_core(hass: HomeAssistant, supervisor_client: AsyncMock) -> None:
@@ -746,6 +860,43 @@ async def test_update_os_with_error(
         )
 
 
+async def test_update_os_with_backup_and_error(
+    hass: HomeAssistant,
+    supervisor_client: AsyncMock,
+) -> None:
+    """Test updating OS update entity with error."""
+    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 setup_backup_integration(hass)
+
+    supervisor_client.os.update.return_value = None
+    supervisor_client.mounts.info.return_value.default_backup_mount = None
+    with (
+        patch(
+            "homeassistant.components.backup.manager.BackupManager.async_create_backup",
+            side_effect=BackupManagerError,
+        ),
+        pytest.raises(HomeAssistantError, match=r"^Error creating backup:"),
+    ):
+        await hass.services.async_call(
+            "update",
+            "install",
+            {
+                "entity_id": "update.home_assistant_operating_system_update",
+                "backup": True,
+            },
+            blocking=True,
+        )
+
+
 async def test_update_supervisor_with_error(
     hass: HomeAssistant, supervisor_client: AsyncMock
 ) -> None: