From e454481d752cb2cb35c82ae80977b7ebea539892 Mon Sep 17 00:00:00 2001
From: Mike Degatano <michael.degatano@gmail.com>
Date: Tue, 17 Dec 2024 03:55:04 +0000
Subject: [PATCH] Add restore from location option

---
 aiohasupervisor/models/backups.py |  1 +
 tests/test_backups.py             | 35 +++++++++++++++++++++++++------
 2 files changed, 30 insertions(+), 6 deletions(-)

diff --git a/aiohasupervisor/models/backups.py b/aiohasupervisor/models/backups.py
index e00517b..7c5314c 100644
--- a/aiohasupervisor/models/backups.py
+++ b/aiohasupervisor/models/backups.py
@@ -174,6 +174,7 @@ class FullRestoreOptions(Request):
 
     password: str | None = None
     background: bool | None = None
+    location: str | None = DEFAULT  # type: ignore[assignment]
 
 
 @dataclass(frozen=True, slots=True)
diff --git a/tests/test_backups.py b/tests/test_backups.py
index fcde508..67d606e 100644
--- a/tests/test_backups.py
+++ b/tests/test_backups.py
@@ -17,6 +17,7 @@
     Folder,
     FreezeOptions,
     FullBackupOptions,
+    FullRestoreOptions,
     PartialBackupOptions,
     PartialRestoreOptions,
     RemoveBackupOptions,
@@ -372,8 +373,20 @@ async def test_remove_backup(
     }
 
 
+@pytest.mark.parametrize(
+    "options",
+    [
+        None,
+        FullRestoreOptions(password="abc123"),  # noqa: S106
+        FullRestoreOptions(background=True),
+        FullRestoreOptions(location=None),
+        FullRestoreOptions(location="test"),
+    ],
+)
 async def test_full_restore(
-    responses: aioresponses, supervisor_client: SupervisorClient
+    responses: aioresponses,
+    supervisor_client: SupervisorClient,
+    options: FullRestoreOptions | None,
 ) -> None:
     """Test full restore API."""
     responses.post(
@@ -381,12 +394,24 @@ async def test_full_restore(
         status=200,
         body=load_fixture("backup_restore.json"),
     )
-    result = await supervisor_client.backups.full_restore("abc123")
+    result = await supervisor_client.backups.full_restore("abc123", options)
     assert result.job_id == "dc9dbc16f6ad4de592ffa72c807ca2bf"
 
 
+@pytest.mark.parametrize(
+    "options",
+    [
+        PartialRestoreOptions(addons={"core_ssh"}),
+        PartialRestoreOptions(homeassistant=True, location=None),
+        PartialRestoreOptions(folders={Folder.SHARE, Folder.SSL}, location="test"),
+        PartialRestoreOptions(addons={"core_ssh"}, background=True),
+        PartialRestoreOptions(addons={"core_ssh"}, password="abc123"),  # noqa: S106
+    ],
+)
 async def test_partial_restore(
-    responses: aioresponses, supervisor_client: SupervisorClient
+    responses: aioresponses,
+    supervisor_client: SupervisorClient,
+    options: PartialRestoreOptions,
 ) -> None:
     """Test partial restore API."""
     responses.post(
@@ -394,9 +419,7 @@ async def test_partial_restore(
         status=200,
         body=load_fixture("backup_restore.json"),
     )
-    result = await supervisor_client.backups.partial_restore(
-        "abc123", PartialRestoreOptions(addons={"core_ssh"})
-    )
+    result = await supervisor_client.backups.partial_restore("abc123", options)
     assert result.job_id == "dc9dbc16f6ad4de592ffa72c807ca2bf"