From 760168de832779e5cdbdab34d1485167ee946644 Mon Sep 17 00:00:00 2001 From: Erik Montnemery Date: Mon, 20 Jan 2025 12:58:17 +0100 Subject: [PATCH] Allow backup writer to update progress during restore (#135975) * Allow backup writer to update progress during restore * Clarify comment --- homeassistant/components/backup/__init__.py | 2 ++ homeassistant/components/backup/manager.py | 13 ++++++++++--- homeassistant/components/hassio/backup.py | 10 ++++++---- tests/components/backup/test_manager.py | 9 +++++++++ 4 files changed, 27 insertions(+), 7 deletions(-) diff --git a/homeassistant/components/backup/__init__.py b/homeassistant/components/backup/__init__.py index f3fe2246ad1..93cadcfb2f3 100644 --- a/homeassistant/components/backup/__init__.py +++ b/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", ] diff --git a/homeassistant/components/backup/manager.py b/homeassistant/components/backup/manager.py index 58600d0a4c0..32979194980 100644 --- a/homeassistant/components/backup/manager.py +++ b/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) diff --git a/homeassistant/components/hassio/backup.py b/homeassistant/components/hassio/backup.py index 537588e856a..76196cfe9e8 100644 --- a/homeassistant/components/hassio/backup.py +++ b/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() diff --git a/tests/components/backup/test_manager.py b/tests/components/backup/test_manager.py index 224f87bea47..4c7eaf634b3 100644 --- a/tests/components/backup/test_manager.py +++ b/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,