Persist hassio backup restore status after core restart (#136857)

* Persist hassio backup restore status after core restart

* Remove useless condition

homeassistant/components/backup/__init__.py

@@ -31,6 +31,7 @@ from .manager import (
     ManagerBackup,
     NewBackup,
     RestoreBackupEvent,
+    RestoreBackupState,
     WrittenBackup,
 )
 from .models import AddonInfo, AgentBackup, Folder
@@ -54,6 +55,7 @@ __all__ = [
     "ManagerBackup",
     "NewBackup",
     "RestoreBackupEvent",
+    "RestoreBackupState",
     "WrittenBackup",
     "async_get_manager",
 ]

homeassistant/components/hassio/backup.py

@@ -5,8 +5,10 @@ from __future__ import annotations
 import asyncio
 from collections.abc import AsyncIterator, Callable, Coroutine, Mapping
 import logging
+import os
 from pathlib import Path
 from typing import Any, cast
+from uuid import UUID
 
 from aiohasupervisor import SupervisorClient
 from aiohasupervisor.exceptions import (
@@ -33,6 +35,7 @@ from homeassistant.components.backup import (
     IncorrectPasswordError,
     NewBackup,
     RestoreBackupEvent,
+    RestoreBackupState,
     WrittenBackup,
     async_get_manager as async_get_backup_manager,
 )
@@ -47,6 +50,7 @@ from .handler import get_supervisor_client
 LOCATION_CLOUD_BACKUP = ".cloud_backup"
 LOCATION_LOCAL = ".local"
 MOUNT_JOBS = ("mount_manager_create_mount", "mount_manager_remove_mount")
+RESTORE_JOB_ID_ENV = "SUPERVISOR_RESTORE_JOB_ID"
 _LOGGER = logging.getLogger(__name__)
 
 
@@ -518,6 +522,37 @@ class SupervisorBackupReaderWriter(BackupReaderWriter):
         on_progress: Callable[[RestoreBackupEvent | IdleEvent], None],
     ) -> None:
         """Check restore status after core restart."""
+        if not (restore_job_id := os.environ.get(RESTORE_JOB_ID_ENV)):
+            _LOGGER.debug("No restore job ID found in environment")
+            return
+
+        _LOGGER.debug("Found restore job ID %s in environment", restore_job_id)
+
+        @callback
+        def on_job_progress(data: Mapping[str, Any]) -> None:
+            """Handle backup restore progress."""
+            if data.get("done") is not True:
+                on_progress(
+                    RestoreBackupEvent(
+                        reason="", stage=None, state=RestoreBackupState.IN_PROGRESS
+                    )
+                )
+                return
+
+            on_progress(
+                RestoreBackupEvent(
+                    reason="", stage=None, state=RestoreBackupState.COMPLETED
+                )
+            )
+            on_progress(IdleEvent())
+            unsub()
+
+        unsub = self._async_listen_job_events(restore_job_id, on_job_progress)
+        try:
+            await self._get_job_state(restore_job_id, on_job_progress)
+        except SupervisorError as err:
+            _LOGGER.debug("Could not get restore job %s: %s", restore_job_id, err)
+            unsub()
 
     @callback
     def _async_listen_job_events(
@@ -546,6 +581,14 @@ class SupervisorBackupReaderWriter(BackupReaderWriter):
         )
         return unsub
 
+    async def _get_job_state(
+        self, job_id: str, on_event: Callable[[Mapping[str, Any]], None]
+    ) -> None:
+        """Poll a job for its state."""
+        job = await self._client.jobs.get_job(UUID(job_id))
+        _LOGGER.debug("Job state: %s", job)
+        on_event(job.to_dict())
+
 
 async def _default_agent(client: SupervisorClient) -> str:
     """Return the default agent for creating a backup."""

tests/components/conftest.py

@@ -535,6 +535,7 @@ def supervisor_client() -> Generator[AsyncMock]:
     supervisor_client.discovery = AsyncMock()
     supervisor_client.homeassistant = AsyncMock()
     supervisor_client.host = AsyncMock()
+    supervisor_client.jobs = AsyncMock()
     supervisor_client.mounts.info.return_value = mounts_info_mock
     supervisor_client.os = AsyncMock()
     supervisor_client.resolution = AsyncMock()

tests/components/hassio/test_backup.py

@@ -13,6 +13,7 @@ from io import StringIO
 import os
 from typing import Any
 from unittest.mock import ANY, AsyncMock, Mock, patch
+from uuid import UUID
 
 from aiohasupervisor.exceptions import (
     SupervisorBadRequestError,
@@ -21,6 +22,7 @@ from aiohasupervisor.exceptions import (
 )
 from aiohasupervisor.models import (
     backups as supervisor_backups,
+    jobs as supervisor_jobs,
     mounts as supervisor_mounts,
 )
 from aiohasupervisor.models.mounts import MountsInfo
@@ -35,7 +37,11 @@ from homeassistant.components.backup import (
     Folder,
 )
 from homeassistant.components.hassio import DOMAIN
-from homeassistant.components.hassio.backup import LOCATION_CLOUD_BACKUP, LOCATION_LOCAL
+from homeassistant.components.hassio.backup import (
+    LOCATION_CLOUD_BACKUP,
+    LOCATION_LOCAL,
+    RESTORE_JOB_ID_ENV,
+)
 from homeassistant.core import HomeAssistant
 from homeassistant.setup import async_setup_component
 
@@ -1802,3 +1808,69 @@ async def test_reader_writer_restore_wrong_parameters(
         "code": "home_assistant_error",
         "message": expected_error,
     }
+
+
+@pytest.mark.usefixtures("hassio_client")
+async def test_restore_progress_after_restart(
+    hass: HomeAssistant,
+    hass_ws_client: WebSocketGenerator,
+    supervisor_client: AsyncMock,
+) -> None:
+    """Test restore backup progress after restart."""
+
+    supervisor_client.jobs.get_job.return_value = supervisor_jobs.Job(
+        name="backup_manager_partial_backup",
+        reference="1ef41507",
+        uuid=UUID("d17bd02be1f0437fa7264b16d38f700e"),
+        progress=0.0,
+        stage="copy_additional_locations",
+        done=True,
+        errors=[],
+        child_jobs=[],
+    )
+
+    with patch.dict(
+        os.environ,
+        MOCK_ENVIRON | {RESTORE_JOB_ID_ENV: "d17bd02be1f0437fa7264b16d38f700e"},
+    ):
+        assert await async_setup_component(hass, BACKUP_DOMAIN, {BACKUP_DOMAIN: {}})
+
+    client = await hass_ws_client(hass)
+
+    await client.send_json_auto_id({"type": "backup/info"})
+    response = await client.receive_json()
+
+    assert response["success"]
+    assert response["result"]["last_non_idle_event"] == {
+        "manager_state": "restore_backup",
+        "reason": "",
+        "stage": None,
+        "state": "completed",
+    }
+    assert response["result"]["state"] == "idle"
+
+
+@pytest.mark.usefixtures("hassio_client")
+async def test_restore_progress_after_restart_unknown_job(
+    hass: HomeAssistant,
+    hass_ws_client: WebSocketGenerator,
+    supervisor_client: AsyncMock,
+) -> None:
+    """Test restore backup progress after restart."""
+
+    supervisor_client.jobs.get_job.side_effect = SupervisorError
+
+    with patch.dict(
+        os.environ,
+        MOCK_ENVIRON | {RESTORE_JOB_ID_ENV: "d17bd02be1f0437fa7264b16d38f700e"},
+    ):
+        assert await async_setup_component(hass, BACKUP_DOMAIN, {BACKUP_DOMAIN: {}})
+
+    client = await hass_ws_client(hass)
+
+    await client.send_json_auto_id({"type": "backup/info"})
+    response = await client.receive_json()
+
+    assert response["success"]
+    assert response["result"]["last_non_idle_event"] is None
+    assert response["result"]["state"] == "idle"