Minor improvement of hassio backup tests (#139775)

tests/components/hassio/test_backup.py

@@ -6,6 +6,7 @@ from collections.abc import (
     Callable,
     Coroutine,
     Generator,
+    Iterable,
 )
 from dataclasses import replace
 from datetime import datetime
@@ -38,6 +39,7 @@ from homeassistant.components.backup import (
     AgentBackup,
     BackupAgent,
     BackupAgentPlatformProtocol,
+    BackupNotFound,
     Folder,
     store as backup_store,
 )
@@ -326,43 +328,70 @@ async def setup_backup_integration(
     await hass.async_block_till_done()
 
 
-class BackupAgentTest(BackupAgent):
+async def aiter_from_iter(iterable: Iterable) -> AsyncIterator:
-    """Test backup agent."""
+    """Convert an iterable to an async iterator."""
+    for i in iterable:
+        yield i
 
-    def __init__(self, name: str, domain: str = "test") -> None:
-        """Initialize the backup agent."""
-        self.domain = domain
-        self.name = name
-        self.unique_id = name
 
-    async def async_download_backup(
+def mock_backup_agent(
-        self, backup_id: str, **kwargs: Any
+    name: str, domain: str = "test", backups: list[AgentBackup] | None = None
-    ) -> AsyncIterator[bytes]:
+) -> Mock:
-        """Download a backup file."""
+    """Create a mock backup agent."""
-        return AsyncMock(spec_set=["__aiter__"])
 
-    async def async_upload_backup(
+    async def delete_backup(backup_id: str, **kwargs: Any) -> None:
-        self,
+        """Mock delete."""
+        get_backup(backup_id)
+
+    async def download_backup(backup_id: str, **kwargs: Any) -> AsyncIterator[bytes]:
+        """Mock download."""
+        return aiter_from_iter((backups_data.get(backup_id, b"backup data"),))
+
+    async def get_backup(backup_id: str, **kwargs: Any) -> AgentBackup:
+        """Get a backup."""
+        backup = next((b for b in backups if b.backup_id == backup_id), None)
+        if backup is None:
+            raise BackupNotFound
+        return backup
+
+    async def upload_backup(
         *,
         open_stream: Callable[[], Coroutine[Any, Any, AsyncIterator[bytes]]],
         backup: AgentBackup,
         **kwargs: Any,
     ) -> None:
         """Upload a backup."""
-        await open_stream()
+        backups.append(backup)
+        backup_stream = await open_stream()
+        backup_data = bytearray()
+        async for chunk in backup_stream:
+            backup_data += chunk
+        backups_data[backup.backup_id] = backup_data
 
-    async def async_list_backups(self, **kwargs: Any) -> list[AgentBackup]:
+    backups = backups or []
-        """List backups."""
+    backups_data: dict[str, bytes] = {}
-        return []
+    mock_agent = Mock(spec=BackupAgent)
-
+    mock_agent.domain = domain
-    async def async_get_backup(
+    mock_agent.name = name
-        self, backup_id: str, **kwargs: Any
+    mock_agent.unique_id = name
-    ) -> AgentBackup | None:
+    type(mock_agent).agent_id = BackupAgent.agent_id
-        """Return a backup."""
+    mock_agent.async_delete_backup = AsyncMock(
-        return None
+        side_effect=delete_backup, spec_set=[BackupAgent.async_delete_backup]
-
+    )
-    async def async_delete_backup(self, backup_id: str, **kwargs: Any) -> None:
+    mock_agent.async_download_backup = AsyncMock(
-        """Delete a backup file."""
+        side_effect=download_backup, spec_set=[BackupAgent.async_download_backup]
+    )
+    mock_agent.async_get_backup = AsyncMock(
+        side_effect=get_backup, spec_set=[BackupAgent.async_get_backup]
+    )
+    mock_agent.async_list_backups = AsyncMock(
+        return_value=backups, spec_set=[BackupAgent.async_list_backups]
+    )
+    mock_agent.async_upload_backup = AsyncMock(
+        side_effect=upload_backup,
+        spec_set=[BackupAgent.async_upload_backup],
+    )
+    return mock_agent
 
 
 async def _setup_backup_platform(
@@ -383,7 +412,7 @@ async def _setup_backup_platform(
     [
         (
             MountsInfo(default_backup_mount=None, mounts=[]),
-            [BackupAgentTest("local", DOMAIN)],
+            [mock_backup_agent("local", DOMAIN)],
         ),
         (
             MountsInfo(
@@ -401,7 +430,7 @@ async def _setup_backup_platform(
                     )
                 ],
             ),
-            [BackupAgentTest("local", DOMAIN), BackupAgentTest("test", DOMAIN)],
+            [mock_backup_agent("local", DOMAIN), mock_backup_agent("test", DOMAIN)],
         ),
         (
             MountsInfo(
@@ -419,7 +448,7 @@ async def _setup_backup_platform(
                     )
                 ],
             ),
-            [BackupAgentTest("local", DOMAIN)],
+            [mock_backup_agent("local", DOMAIN)],
         ),
     ],
 )
@@ -576,36 +605,9 @@ async def test_agent_upload(
 ) -> None:
     """Test agent upload backup."""
     client = await hass_client()
-    backup_id = "test-backup"
     supervisor_client.backups.backup_info.return_value = TEST_BACKUP_DETAILS
-    test_backup = AgentBackup(
-        addons=[AddonInfo(name="Test", slug="test", version="1.0.0")],
-        backup_id=backup_id,
-        database_included=True,
-        date="1970-01-01T00:00:00.000Z",
-        extra_metadata={},
-        folders=[Folder.MEDIA, Folder.SHARE],
-        homeassistant_included=True,
-        homeassistant_version="2024.12.0",
-        name="Test",
-        protected=False,
-        size=0,
-    )
 
     supervisor_client.backups.reload.assert_not_called()
-    with (
-        patch("pathlib.Path.mkdir"),
-        patch("pathlib.Path.open"),
-        patch(
-            "homeassistant.components.backup.manager.BackupManager.async_get_backup",
-        ) as fetch_backup,
-        patch(
-            "homeassistant.components.backup.manager.read_backup",
-            return_value=test_backup,
-        ),
-        patch("shutil.copy"),
-    ):
-        fetch_backup.return_value = test_backup
     resp = await client.post(
         "/api/backup/upload?agent_id=hassio.local",
         data={"file": StringIO("test")},
@@ -1551,7 +1553,7 @@ async def test_reader_writer_create_download_remove_error(
     method_mock = getattr(supervisor_client.backups, method)
     method_mock.side_effect = exception
 
-    remote_agent = BackupAgentTest("remote")
+    remote_agent = mock_backup_agent("remote")
     await _setup_backup_platform(
         hass,
         domain="test",
@@ -1636,7 +1638,7 @@ async def test_reader_writer_create_info_error(
     supervisor_client.backups.backup_info.side_effect = exception
     supervisor_client.jobs.get_job.return_value = TEST_JOB_NOT_DONE
 
-    remote_agent = BackupAgentTest("remote")
+    remote_agent = mock_backup_agent("remote")
     await _setup_backup_platform(
         hass,
         domain="test",
@@ -1713,7 +1715,7 @@ async def test_reader_writer_create_remote_backup(
     supervisor_client.backups.backup_info.return_value = TEST_BACKUP_DETAILS_5
     supervisor_client.jobs.get_job.return_value = TEST_JOB_NOT_DONE
 
-    remote_agent = BackupAgentTest("remote")
+    remote_agent = mock_backup_agent("remote")
     await _setup_backup_platform(
         hass,
         domain="test",
@@ -1861,24 +1863,10 @@ async def test_agent_receive_remote_backup(
 ) -> None:
     """Test receiving a backup which will be uploaded to a remote agent."""
     client = await hass_client()
-    backup_id = "test-backup"
     supervisor_client.backups.backup_info.return_value = TEST_BACKUP_DETAILS_5
     supervisor_client.backups.upload_backup.return_value = "test_slug"
-    test_backup = AgentBackup(
-        addons=[AddonInfo(name="Test", slug="test", version="1.0.0")],
-        backup_id=backup_id,
-        database_included=True,
-        date="1970-01-01T00:00:00.000Z",
-        extra_metadata={},
-        folders=[Folder.MEDIA, Folder.SHARE],
-        homeassistant_included=True,
-        homeassistant_version="2024.12.0",
-        name="Test",
-        protected=False,
-        size=0.0,
-    )
 
-    remote_agent = BackupAgentTest("remote")
+    remote_agent = mock_backup_agent("remote")
     await _setup_backup_platform(
         hass,
         domain="test",
@@ -1889,19 +1877,6 @@ async def test_agent_receive_remote_backup(
     )
 
     supervisor_client.backups.reload.assert_not_called()
-    with (
-        patch("pathlib.Path.mkdir"),
-        patch("pathlib.Path.open"),
-        patch(
-            "homeassistant.components.backup.manager.BackupManager.async_get_backup",
-        ) as fetch_backup,
-        patch(
-            "homeassistant.components.backup.manager.read_backup",
-            return_value=test_backup,
-        ),
-        patch("shutil.copy"),
-    ):
-        fetch_backup.return_value = test_backup
     resp = await client.post(
         "/api/backup/upload?agent_id=test.remote",
         data={"file": StringIO("test")},
@@ -1996,6 +1971,103 @@ async def test_reader_writer_restore(
     assert response["result"] is None
 
 
+@pytest.mark.usefixtures("hassio_client", "setup_backup_integration")
+async def test_reader_writer_restore_remote_backup(
+    hass: HomeAssistant,
+    hass_ws_client: WebSocketGenerator,
+    supervisor_client: AsyncMock,
+) -> None:
+    """Test restoring a backup from a remote agent."""
+    client = await hass_ws_client(hass)
+    supervisor_client.backups.partial_restore.return_value.job_id = UUID(TEST_JOB_ID)
+    supervisor_client.backups.list.return_value = [TEST_BACKUP_5]
+    supervisor_client.backups.backup_info.return_value = TEST_BACKUP_DETAILS_5
+    supervisor_client.jobs.get_job.return_value = TEST_JOB_NOT_DONE
+
+    backup_id = "abc123"
+    test_backup = AgentBackup(
+        addons=[AddonInfo(name="Test", slug="test", version="1.0.0")],
+        backup_id=backup_id,
+        database_included=True,
+        date="1970-01-01T00:00:00.000Z",
+        extra_metadata={},
+        folders=[Folder.MEDIA, Folder.SHARE],
+        homeassistant_included=True,
+        homeassistant_version="2024.12.0",
+        name="Test",
+        protected=False,
+        size=0.0,
+    )
+    remote_agent = mock_backup_agent("remote", backups=[test_backup])
+    await _setup_backup_platform(
+        hass,
+        domain="test",
+        platform=Mock(
+            async_get_backup_agents=AsyncMock(return_value=[remote_agent]),
+            spec_set=BackupAgentPlatformProtocol,
+        ),
+    )
+
+    await client.send_json_auto_id({"type": "backup/subscribe_events"})
+    response = await client.receive_json()
+    assert response["event"] == {
+        "manager_state": "idle",
+    }
+    response = await client.receive_json()
+    assert response["success"]
+
+    await client.send_json_auto_id(
+        {"type": "backup/restore", "agent_id": "test.remote", "backup_id": backup_id}
+    )
+    response = await client.receive_json()
+    assert response["event"] == {
+        "manager_state": "restore_backup",
+        "reason": None,
+        "stage": None,
+        "state": "in_progress",
+    }
+
+    remote_agent.async_download_backup.assert_called_once_with(backup_id)
+    assert len(remote_agent.async_get_backup.mock_calls) == 2
+    for call in remote_agent.async_get_backup.mock_calls:
+        assert call.args[0] == backup_id
+    supervisor_client.backups.partial_restore.assert_called_once_with(
+        backup_id,
+        supervisor_backups.PartialRestoreOptions(
+            addons=None,
+            background=True,
+            folders=None,
+            homeassistant=True,
+            location=LOCATION_CLOUD_BACKUP,
+            password=None,
+        ),
+    )
+
+    await client.send_json_auto_id(
+        {
+            "type": "supervisor/event",
+            "data": {"event": "job", "data": {"done": True, "uuid": TEST_JOB_ID}},
+        }
+    )
+    response = await client.receive_json()
+    assert response["success"]
+
+    response = await client.receive_json()
+    assert response["event"] == {
+        "manager_state": "restore_backup",
+        "reason": None,
+        "stage": None,
+        "state": "completed",
+    }
+
+    response = await client.receive_json()
+    assert response["event"] == {"manager_state": "idle"}
+
+    response = await client.receive_json()
+    assert response["success"]
+    assert response["result"] is None
+
+
 @pytest.mark.usefixtures("hassio_client", "setup_backup_integration")
 async def test_reader_writer_restore_report_progress(
     hass: HomeAssistant,