Adapt hassio backup agent to supervisor changes (#133428)

homeassistant/components/hassio/backup.py

@@ -8,7 +8,10 @@ import logging
 from pathlib import Path
 from typing import Any, cast
 
-from aiohasupervisor.exceptions import SupervisorBadRequestError
+from aiohasupervisor.exceptions import (
+    SupervisorBadRequestError,
+    SupervisorNotFoundError,
+)
 from aiohasupervisor.models import (
     backups as supervisor_backups,
     mounts as supervisor_mounts,
@@ -130,7 +133,10 @@ class SupervisorBackupAgent(BackupAgent):
         **kwargs: Any,
     ) -> AsyncIterator[bytes]:
         """Download a backup file."""
-        return await self._client.backups.download_backup(backup_id)
+        return await self._client.backups.download_backup(
+            backup_id,
+            options=supervisor_backups.DownloadBackupOptions(location=self.location),
+        )
 
     async def async_upload_backup(
         self,
@@ -169,11 +175,18 @@ class SupervisorBackupAgent(BackupAgent):
     async def async_delete_backup(self, backup_id: str, **kwargs: Any) -> None:
         """Remove a backup."""
         try:
-            await self._client.backups.remove_backup(backup_id)
+            await self._client.backups.remove_backup(
+                backup_id,
+                options=supervisor_backups.RemoveBackupOptions(
+                    location={self.location}
+                ),
+            )
         except SupervisorBadRequestError as err:
             if err.args[0] != "Backup does not exist":
                 raise
             _LOGGER.debug("Backup %s does not exist", backup_id)
+        except SupervisorNotFoundError:
+            _LOGGER.debug("Backup %s does not exist", backup_id)
 
 
 class SupervisorBackupReaderWriter(BackupReaderWriter):
@@ -200,7 +213,11 @@ class SupervisorBackupReaderWriter(BackupReaderWriter):
         """Create a backup."""
         manager = self._hass.data[DATA_MANAGER]
 
-        include_addons_set = set(include_addons) if include_addons else None
+        include_addons_set: supervisor_backups.AddonSet | set[str] | None = None
+        if include_all_addons:
+            include_addons_set = supervisor_backups.AddonSet.ALL
+        elif include_addons:
+            include_addons_set = set(include_addons)
         include_folders_set = (
             {supervisor_backups.Folder(folder) for folder in include_folders}
             if include_folders
@@ -266,7 +283,12 @@ class SupervisorBackupReaderWriter(BackupReaderWriter):
         async def remove_backup() -> None:
             if not remove_after_upload:
                 return
-            await self._client.backups.remove_backup(backup_id)
+            await self._client.backups.remove_backup(
+                backup_id,
+                options=supervisor_backups.RemoveBackupOptions(
+                    location={LOCATION_CLOUD_BACKUP}
+                ),
+            )
 
         details = await self._client.backups.backup_info(backup_id)
 
@@ -306,7 +328,12 @@ class SupervisorBackupReaderWriter(BackupReaderWriter):
         async def remove_backup() -> None:
             if locations:
                 return
-            await self._client.backups.remove_backup(backup_id)
+            await self._client.backups.remove_backup(
+                backup_id,
+                options=supervisor_backups.RemoveBackupOptions(
+                    location={LOCATION_CLOUD_BACKUP}
+                ),
+            )
 
         details = await self._client.backups.backup_info(backup_id)
 
@@ -341,6 +368,7 @@ class SupervisorBackupReaderWriter(BackupReaderWriter):
         )
 
         manager = self._hass.data[DATA_MANAGER]
+        restore_location: str | None
         if manager.backup_agents[agent_id].domain != DOMAIN:
             # Download the backup to the supervisor. Supervisor will clean up the backup
             # two days after the restore is done.
@@ -349,6 +377,10 @@ class SupervisorBackupReaderWriter(BackupReaderWriter):
                 stream=await open_stream(),
                 suggested_filename=f"{backup_id}.tar",
             )
+            restore_location = LOCATION_CLOUD_BACKUP
+        else:
+            agent = cast(SupervisorBackupAgent, manager.backup_agents[agent_id])
+            restore_location = agent.location
 
         job = await self._client.backups.partial_restore(
             backup_id,
@@ -358,6 +390,7 @@ class SupervisorBackupReaderWriter(BackupReaderWriter):
                 homeassistant=restore_homeassistant,
                 password=password,
                 background=True,
+                location=restore_location,
             ),
         )
 

tests/components/hassio/test_backup.py

@@ -14,7 +14,10 @@ import os
 from typing import Any
 from unittest.mock import AsyncMock, Mock, patch
 
-from aiohasupervisor.exceptions import SupervisorBadRequestError
+from aiohasupervisor.exceptions import (
+    SupervisorBadRequestError,
+    SupervisorNotFoundError,
+)
 from aiohasupervisor.models import (
     backups as supervisor_backups,
     mounts as supervisor_mounts,
@@ -403,6 +406,10 @@ async def test_agent_download(
     assert resp.status == 200
     assert await resp.content.read() == b"backup data"
 
+    supervisor_client.backups.download_backup.assert_called_once_with(
+        "abc123", options=supervisor_backups.DownloadBackupOptions(location=None)
+    )
+
 
 @pytest.mark.usefixtures("hassio_client", "setup_integration")
 async def test_agent_download_unavailable_backup(
@@ -491,7 +498,9 @@ async def test_agent_delete_backup(
 
     assert response["success"]
     assert response["result"] == {"agent_errors": {}}
-    supervisor_client.backups.remove_backup.assert_called_once_with(backup_id)
+    supervisor_client.backups.remove_backup.assert_called_once_with(
+        backup_id, options=supervisor_backups.RemoveBackupOptions(location={None})
+    )
 
 
 @pytest.mark.usefixtures("hassio_client", "setup_integration")
@@ -512,6 +521,13 @@ async def test_agent_delete_backup(
                 "result": {"agent_errors": {}},
             },
         ),
+        (
+            SupervisorNotFoundError(),
+            {
+                "success": True,
+                "result": {"agent_errors": {}},
+            },
+        ),
     ],
 )
 async def test_agent_delete_with_error(
@@ -535,7 +551,9 @@ async def test_agent_delete_with_error(
     response = await client.receive_json()
 
     assert response == {"id": 1, "type": "result"} | expected_response
-    supervisor_client.backups.remove_backup.assert_called_once_with(backup_id)
+    supervisor_client.backups.remove_backup.assert_called_once_with(
+        backup_id, options=supervisor_backups.RemoveBackupOptions(location={None})
+    )
 
 
 @pytest.mark.usefixtures("hassio_client", "setup_integration")
@@ -627,7 +645,7 @@ DEFAULT_BACKUP_OPTIONS = supervisor_backups.PartialBackupOptions(
         ),
         (
             {"include_all_addons": True},
-            DEFAULT_BACKUP_OPTIONS,
+            replace(DEFAULT_BACKUP_OPTIONS, addons="all"),
         ),
         (
             {"include_database": False},
@@ -782,7 +800,10 @@ async def test_reader_writer_create_remote_backup(
     }
 
     supervisor_client.backups.download_backup.assert_called_once_with("test_slug")
-    supervisor_client.backups.remove_backup.assert_called_once_with("test_slug")
+    supervisor_client.backups.remove_backup.assert_called_once_with(
+        "test_slug",
+        options=supervisor_backups.RemoveBackupOptions({LOCATION_CLOUD_BACKUP}),
+    )
 
 
 @pytest.mark.usefixtures("hassio_client", "setup_integration")
@@ -895,7 +916,10 @@ async def test_agent_receive_remote_backup(
     assert resp.status == 201
 
     supervisor_client.backups.download_backup.assert_called_once_with("test_slug")
-    supervisor_client.backups.remove_backup.assert_called_once_with("test_slug")
+    supervisor_client.backups.remove_backup.assert_called_once_with(
+        "test_slug",
+        options=supervisor_backups.RemoveBackupOptions({LOCATION_CLOUD_BACKUP}),
+    )
 
 
 @pytest.mark.usefixtures("hassio_client", "setup_integration")
@@ -933,6 +957,7 @@ async def test_reader_writer_restore(
             background=True,
             folders=None,
             homeassistant=True,
+            location=None,
             password=None,
         ),
     )