Poll supervisor job state when creating or restoring a backup (#136891)

* Poll supervisor job state when creating or restoring a backup * Update tests * Add tests for create and restore jobs finishing early
home-assistant · Jan 30, 2025 · 0764aca · 0764aca
1 parent 8babdc0
commit 0764aca
Show file tree

Hide file tree

Showing 2 changed files with 136 additions and 52 deletions.
diff --git a/homeassistant/components/hassio/backup.py b/homeassistant/components/hassio/backup.py
@@ -350,8 +350,9 @@ def on_job_progress(data: Mapping[str, Any]) -> None:
                 backup_id = data.get("reference")
                 backup_complete.set()
 
+        unsub = self._async_listen_job_events(backup.job_id, on_job_progress)
         try:
-            unsub = self._async_listen_job_events(backup.job_id, on_job_progress)
+            await self._get_job_state(backup.job_id, on_job_progress)
             await backup_complete.wait()
         finally:
             unsub()
@@ -506,12 +507,13 @@ async def async_restore_backup(
 
         @callback
         def on_job_progress(data: Mapping[str, Any]) -> None:
-            """Handle backup progress."""
+            """Handle backup restore progress."""
             if data.get("done") is True:
                 restore_complete.set()
 
+        unsub = self._async_listen_job_events(job.job_id, on_job_progress)
         try:
-            unsub = self._async_listen_job_events(job.job_id, on_job_progress)
+            await self._get_job_state(job.job_id, on_job_progress)
             await restore_complete.wait()
         finally:
             unsub()