Allow backup writer to update progress during restore (#135975)

* Allow backup writer to update progress during restore

* Clarify comment

homeassistant/components/backup/__init__.py

@@ -27,6 +27,7 @@ from .manager import (
     IncorrectPasswordError,
     ManagerBackup,
     NewBackup,
+    RestoreBackupEvent,
     WrittenBackup,
 )
 from .models import AddonInfo, AgentBackup, Folder
@@ -47,6 +48,7 @@ __all__ = [
     "LocalBackupAgent",
     "ManagerBackup",
     "NewBackup",
+    "RestoreBackupEvent",
     "WrittenBackup",
 ]
 

homeassistant/components/backup/manager.py

@@ -147,6 +147,7 @@ class RestoreBackupState(StrEnum):
     """Receive backup state enum."""
 
     COMPLETED = "completed"
+    CORE_RESTART = "core_restart"
     FAILED = "failed"
     IN_PROGRESS = "in_progress"
 
@@ -217,7 +218,7 @@ class BackupReaderWriter(abc.ABC):
         include_database: bool,
         include_folders: list[Folder] | None,
         include_homeassistant: bool,
-        on_progress: Callable[[ManagerStateEvent], None],
+        on_progress: Callable[[CreateBackupEvent], None],
         password: str | None,
     ) -> tuple[NewBackup, asyncio.Task[WrittenBackup]]:
         """Create a backup."""
@@ -238,6 +239,7 @@ class BackupReaderWriter(abc.ABC):
         backup_id: str,
         *,
         agent_id: str,
+        on_progress: Callable[[RestoreBackupEvent], None],
         open_stream: Callable[[], Coroutine[Any, Any, AsyncIterator[bytes]]],
         password: str | None,
         restore_addons: list[str] | None,
@@ -941,6 +943,7 @@ class BackupManager:
             backup_id=backup_id,
             open_stream=open_backup,
             agent_id=agent_id,
+            on_progress=self.async_on_backup_event,
             password=password,
             restore_addons=restore_addons,
             restore_database=restore_database,
@@ -1130,7 +1133,7 @@ class CoreBackupReaderWriter(BackupReaderWriter):
         include_database: bool,
         include_folders: list[Folder] | None,
         include_homeassistant: bool,
-        on_progress: Callable[[ManagerStateEvent], None],
+        on_progress: Callable[[CreateBackupEvent], None],
         password: str | None,
     ) -> tuple[NewBackup, asyncio.Task[WrittenBackup]]:
         """Initiate generating a backup."""
@@ -1170,7 +1173,7 @@ class CoreBackupReaderWriter(BackupReaderWriter):
         date_str: str,
         extra_metadata: dict[str, bool | str],
         include_database: bool,
-        on_progress: Callable[[ManagerStateEvent], None],
+        on_progress: Callable[[CreateBackupEvent], None],
         password: str | None,
     ) -> WrittenBackup:
         """Generate a backup."""
@@ -1378,6 +1381,7 @@ class CoreBackupReaderWriter(BackupReaderWriter):
         open_stream: Callable[[], Coroutine[Any, Any, AsyncIterator[bytes]]],
         *,
         agent_id: str,
+        on_progress: Callable[[RestoreBackupEvent], None],
         password: str | None,
         restore_addons: list[str] | None,
         restore_database: bool,
@@ -1440,6 +1444,9 @@ class CoreBackupReaderWriter(BackupReaderWriter):
             )
 
         await self._hass.async_add_executor_job(_write_restore_file)
+        on_progress(
+            RestoreBackupEvent(stage=None, state=RestoreBackupState.CORE_RESTART)
+        )
         await self._hass.services.async_call("homeassistant", "restart", blocking=True)
 
 

homeassistant/components/hassio/backup.py

@@ -29,6 +29,7 @@ from homeassistant.components.backup import (
     Folder,
     IncorrectPasswordError,
     NewBackup,
+    RestoreBackupEvent,
     WrittenBackup,
 )
 from homeassistant.core import HomeAssistant, callback
@@ -275,7 +276,7 @@ class SupervisorBackupReaderWriter(BackupReaderWriter):
         backup_id: str | None = None
 
         @callback
-        def on_progress(data: Mapping[str, Any]) -> None:
+        def on_job_progress(data: Mapping[str, Any]) -> None:
             """Handle backup progress."""
             nonlocal backup_id
             if data.get("done") is True:
@@ -283,7 +284,7 @@ class SupervisorBackupReaderWriter(BackupReaderWriter):
                 backup_complete.set()
 
         try:
-            unsub = self._async_listen_job_events(backup.job_id, on_progress)
+            unsub = self._async_listen_job_events(backup.job_id, on_job_progress)
             await backup_complete.wait()
         finally:
             unsub()
@@ -374,6 +375,7 @@ class SupervisorBackupReaderWriter(BackupReaderWriter):
         backup_id: str,
         *,
         agent_id: str,
+        on_progress: Callable[[RestoreBackupEvent], None],
         open_stream: Callable[[], Coroutine[Any, Any, AsyncIterator[bytes]]],
         password: str | None,
         restore_addons: list[str] | None,
@@ -437,13 +439,13 @@ class SupervisorBackupReaderWriter(BackupReaderWriter):
         restore_complete = asyncio.Event()
 
         @callback
-        def on_progress(data: Mapping[str, Any]) -> None:
+        def on_job_progress(data: Mapping[str, Any]) -> None:
             """Handle backup progress."""
             if data.get("done") is True:
                 restore_complete.set()
 
         try:
-            unsub = self._async_listen_job_events(job.job_id, on_progress)
+            unsub = self._async_listen_job_events(job.job_id, on_job_progress)
             await restore_complete.wait()
         finally:
             unsub()

tests/components/backup/test_manager.py

@@ -2302,6 +2302,15 @@ async def test_restore_backup(
             "state": RestoreBackupState.IN_PROGRESS,
         }
 
+        result = await ws_client.receive_json()
+        assert result["event"] == {
+            "manager_state": BackupManagerState.RESTORE_BACKUP,
+            "stage": None,
+            "state": RestoreBackupState.CORE_RESTART,
+        }
+
+        # Note: The core restart is not tested here, in reality the following events
+        # are not sent because the core restart closes the WS connection.
         result = await ws_client.receive_json()
         assert result["event"] == {
             "manager_state": BackupManagerState.RESTORE_BACKUP,