Add remove and download location (#35)

Author: mdegat01

aiohasupervisor/backups.py

@@ -14,12 +14,14 @@
     BackupList,
     BackupsInfo,
     BackupsOptions,
+    DownloadBackupOptions,
     FreezeOptions,
     FullBackupOptions,
     FullRestoreOptions,
     NewBackup,
     PartialBackupOptions,
     PartialRestoreOptions,
+    RemoveBackupOptions,
     UploadBackupOptions,
     UploadedBackup,
 )
@@ -81,9 +83,13 @@ async def backup_info(self, backup: str) -> BackupComplete:
         result = await self._client.get(f"backups/{backup}/info")
         return BackupComplete.from_dict(result.data)
 
-    async def remove_backup(self, backup: str) -> None:
+    async def remove_backup(
+        self, backup: str, options: RemoveBackupOptions | None = None
+    ) -> None:
         """Remove a backup."""
-        await self._client.delete(f"backups/{backup}")
+        await self._client.delete(
+            f"backups/{backup}", json=options.to_dict() if options else None
+        )
 
     async def full_restore(
         self, backup: str, options: FullRestoreOptions | None = None
@@ -129,9 +135,17 @@ async def upload_backup(
 
         return UploadedBackup.from_dict(result.data).slug
 
-    async def download_backup(self, backup: str) -> AsyncIterator[bytes]:
+    async def download_backup(
+        self, backup: str, options: DownloadBackupOptions | None = None
+    ) -> AsyncIterator[bytes]:
         """Download backup and return stream."""
+        params = MultiDict()
+        if options and options.location:
+            params.add("location", options.location or "")
+
         result = await self._client.get(
-            f"backups/{backup}/download", response_type=ResponseType.STREAM
+            f"backups/{backup}/download",
+            params=params,
+            response_type=ResponseType.STREAM,
         )
         return result.data

aiohasupervisor/client.py

@@ -212,6 +212,7 @@ async def delete(
         uri: str,
         *,
         params: dict[str, str] | MultiDict[str] | None = None,
+        json: dict[str, Any] | None = None,
         timeout: ClientTimeout | None = DEFAULT_TIMEOUT,
     ) -> Response:
         """Handle a DELETE request to Supervisor."""
@@ -220,6 +221,7 @@ async def delete(
             uri,
             params=params,
             response_type=ResponseType.NONE,
+            json=json,
             timeout=timeout,
         )
 

aiohasupervisor/models/__init__.py

@@ -33,13 +33,15 @@
     BackupsInfo,
     BackupsOptions,
     BackupType,
+    DownloadBackupOptions,
     Folder,
     FreezeOptions,
     FullBackupOptions,
     FullRestoreOptions,
     NewBackup,
     PartialBackupOptions,
     PartialRestoreOptions,
+    RemoveBackupOptions,
     UploadBackupOptions,
 )
 from aiohasupervisor.models.discovery import (
@@ -209,13 +211,15 @@
     "BackupsInfo",
     "BackupsOptions",
     "BackupType",
+    "DownloadBackupOptions",
     "Folder",
     "FreezeOptions",
     "FullBackupOptions",
     "FullRestoreOptions",
     "NewBackup",
     "PartialBackupOptions",
     "PartialRestoreOptions",
+    "RemoveBackupOptions",
     "UploadBackupOptions",
     "Discovery",
     "DiscoveryConfig",

aiohasupervisor/models/backups.py

@@ -5,7 +5,7 @@
 from datetime import datetime
 from enum import StrEnum
 
-from .base import Request, ResponseData
+from .base import DEFAULT, Request, ResponseData
 
 # --- ENUMS ----
 
@@ -135,7 +135,7 @@ class FullBackupOptions(Request):
     name: str | None = None
     password: str | None = None
     compressed: bool | None = None
-    location: set[str | None] | str | None = None
+    location: list[str | None] | str | None = DEFAULT  # type: ignore[assignment]
     homeassistant_exclude_database: bool | None = None
     background: bool | None = None
     extra: dict | None = None
@@ -185,3 +185,17 @@ class UploadedBackup(ResponseData):
     """UploadedBackup model."""
 
     slug: str
+
+
+@dataclass(frozen=True, slots=True)
+class RemoveBackupOptions(Request):
+    """RemoveBackupOptions model."""
+
+    location: set[str | None] = None
+
+
+@dataclass(frozen=True, slots=True)
+class DownloadBackupOptions(Request):
+    """DownloadBackupOptions model."""
+
+    location: str | None = DEFAULT  # type: ignore[assignment]

tests/test_backups.py

@@ -12,11 +12,13 @@
 from aiohasupervisor import SupervisorClient
 from aiohasupervisor.models import (
     BackupsOptions,
+    DownloadBackupOptions,
     Folder,
     FreezeOptions,
     FullBackupOptions,
     PartialBackupOptions,
     PartialRestoreOptions,
+    RemoveBackupOptions,
     UploadBackupOptions,
 )
 
@@ -131,6 +133,32 @@ async def test_partial_restore_options() -> None:
         PartialRestoreOptions(background=True)
 
 
+async def test_backup_options_location() -> None:
+    """Test location field in backup options."""
+    assert FullBackupOptions(location=["test", None]).to_dict() == {
+        "location": ["test", None]
+    }
+    assert FullBackupOptions(location="test").to_dict() == {"location": "test"}
+    assert FullBackupOptions(location=None).to_dict() == {"location": None}
+    assert FullBackupOptions().to_dict() == {}
+
+    assert PartialBackupOptions(
+        location=["test", None], folders={Folder.SSL}
+    ).to_dict() == {
+        "location": ["test", None],
+        "folders": ["ssl"],
+    }
+    assert PartialBackupOptions(location="test", folders={Folder.SSL}).to_dict() == {
+        "location": "test",
+        "folders": ["ssl"],
+    }
+    assert PartialBackupOptions(location=None, folders={Folder.SSL}).to_dict() == {
+        "location": None,
+        "folders": ["ssl"],
+    }
+    assert PartialBackupOptions(folders={Folder.SSL}).to_dict() == {"folders": ["ssl"]}
+
+
 def backup_callback(url: str, **kwargs: dict[str, Any]) -> CallbackResult:  # noqa: ARG001
     """Return response based on whether backup was in background or not."""
     if kwargs["json"] and kwargs["json"].get("background"):
@@ -323,12 +351,17 @@ async def test_backup_info_with_multiple_locations(
     assert result.locations == {None, "Test"}
 
 
+@pytest.mark.parametrize(
+    "options", [None, RemoveBackupOptions(location={"test", None})]
+)
 async def test_remove_backup(
-    responses: aioresponses, supervisor_client: SupervisorClient
+    responses: aioresponses,
+    supervisor_client: SupervisorClient,
+    options: RemoveBackupOptions | None,
 ) -> None:
     """Test remove backup API."""
     responses.delete(f"{SUPERVISOR_URL}/backups/abc123", status=200)
-    assert await supervisor_client.backups.remove_backup("abc123") is None
+    assert await supervisor_client.backups.remove_backup("abc123", options) is None
     assert responses.requests.keys() == {
         ("DELETE", URL(f"{SUPERVISOR_URL}/backups/abc123"))
     }
@@ -398,14 +431,27 @@ async def test_upload_backup_to_locations(
     assert result == "7fed74c8"
 
 
+@pytest.mark.parametrize(
+    ("options", "query"),
+    [
+        (None, ""),
+        (DownloadBackupOptions(location="test"), "?location=test"),
+        (DownloadBackupOptions(location=None), "?location="),
+    ],
+)
 async def test_download_backup(
-    responses: aioresponses, supervisor_client: SupervisorClient
+    responses: aioresponses,
+    supervisor_client: SupervisorClient,
+    options: DownloadBackupOptions | None,
+    query: str,
 ) -> None:
     """Test download backup API."""
     responses.get(
-        f"{SUPERVISOR_URL}/backups/7fed74c8/download", status=200, body=b"backup test"
+        f"{SUPERVISOR_URL}/backups/7fed74c8/download{query}",
+        status=200,
+        body=b"backup test",
     )
-    result = await supervisor_client.backups.download_backup("7fed74c8")
+    result = await supervisor_client.backups.download_backup("7fed74c8", options)
     assert isinstance(result, AsyncIterator)
     async for chunk in result:
         assert chunk == b"backup test"