Add pagination support for IDrive e2 (home-assistant#162960)

Author: patrickvorgers

homeassistant/components/idrive_e2/backup.py

@@ -329,14 +329,14 @@ async def _list_backups(self) -> dict[str, AgentBackup]:
             return self._backup_cache
 
         backups = {}
-        response = await cast(Any, self._client).list_objects_v2(Bucket=self._bucket)
-
-        # Filter for metadata files only
-        metadata_files = [
-            obj
-            for obj in response.get("Contents", [])
-            if obj["Key"].endswith(".metadata.json")
-        ]
+        paginator = self._client.get_paginator("list_objects_v2")
+        metadata_files: list[dict[str, Any]] = []
+        async for page in paginator.paginate(Bucket=self._bucket):
+            metadata_files.extend(
+                obj
+                for obj in page.get("Contents", [])
+                if obj["Key"].endswith(".metadata.json")
+            )
 
         for metadata_file in metadata_files:
             try:

tests/components/idrive_e2/conftest.py

@@ -4,7 +4,7 @@
 
 from collections.abc import AsyncIterator, Generator
 import json
-from unittest.mock import AsyncMock, patch
+from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
 
@@ -53,9 +53,11 @@ def mock_client(agent_backup: AgentBackup) -> Generator[AsyncMock]:
         client = create_client.return_value
 
         tar_file, metadata_file = suggested_filenames(agent_backup)
-        client.list_objects_v2.return_value = {
-            "Contents": [{"Key": tar_file}, {"Key": metadata_file}]
-        }
+        # Mock the paginator for list_objects_v2
+        client.get_paginator = MagicMock()
+        client.get_paginator.return_value.paginate.return_value.__aiter__.return_value = [
+            {"Contents": [{"Key": tar_file}, {"Key": metadata_file}]}
+        ]
         client.create_multipart_upload.return_value = {"UploadId": "upload_id"}
         client.upload_part.return_value = {"ETag": "etag"}
         client.list_buckets.return_value = {

tests/components/idrive_e2/test_backup.py

@@ -179,7 +179,9 @@ async def test_agents_get_backup_does_not_throw_on_not_found(
     mock_client: MagicMock,
 ) -> None:
     """Test agent get backup does not throw on a backup not found."""
-    mock_client.list_objects_v2.return_value = {"Contents": []}
+    mock_client.get_paginator.return_value.paginate.return_value.__aiter__.return_value = [
+        {"Contents": []}
+    ]
 
     client = await hass_ws_client(hass)
     await client.send_json_auto_id({"type": "backup/details", "backup_id": "random"})
@@ -202,18 +204,20 @@ async def test_agents_list_backups_with_corrupted_metadata(
     agent = IDriveE2BackupAgent(hass, mock_config_entry)
 
     # Set up mock responses for both valid and corrupted metadata files
-    mock_client.list_objects_v2.return_value = {
-        "Contents": [
-            {
-                "Key": "valid_backup.metadata.json",
-                "LastModified": "2023-01-01T00:00:00+00:00",
-            },
-            {
-                "Key": "corrupted_backup.metadata.json",
-                "LastModified": "2023-01-01T00:00:00+00:00",
-            },
-        ]
-    }
+    mock_client.get_paginator.return_value.paginate.return_value.__aiter__.return_value = [
+        {
+            "Contents": [
+                {
+                    "Key": "valid_backup.metadata.json",
+                    "LastModified": "2023-01-01T00:00:00+00:00",
+                },
+                {
+                    "Key": "corrupted_backup.metadata.json",
+                    "LastModified": "2023-01-01T00:00:00+00:00",
+                },
+            ]
+        }
+    ]
 
     # Mock responses for get_object calls
     valid_metadata = json.dumps(agent_backup.as_dict())
@@ -270,7 +274,9 @@ async def test_agents_delete_not_throwing_on_not_found(
     mock_client: MagicMock,
 ) -> None:
     """Test agent delete backup does not throw on a backup not found."""
-    mock_client.list_objects_v2.return_value = {"Contents": []}
+    mock_client.get_paginator.return_value.paginate.return_value.__aiter__.return_value = [
+        {"Contents": []}
+    ]
 
     client = await hass_ws_client(hass)
 
@@ -284,7 +290,7 @@ async def test_agents_delete_not_throwing_on_not_found(
 
     assert response["success"]
     assert response["result"] == {"agent_errors": {}}
-    assert mock_client.delete_object.call_count == 0
+    assert mock_client.delete_objects.call_count == 0
 
 
 async def test_agents_upload(
@@ -490,28 +496,35 @@ async def test_cache_expiration(
     metadata_content = json.dumps(agent_backup.as_dict())
     mock_body = AsyncMock()
     mock_body.read.return_value = metadata_content.encode()
-    mock_client.list_objects_v2.return_value = {
-        "Contents": [
-            {"Key": "test.metadata.json", "LastModified": "2023-01-01T00:00:00+00:00"}
-        ]
-    }
+    mock_client.get_paginator.return_value.paginate.return_value.__aiter__.return_value = [
+        {
+            "Contents": [
+                {
+                    "Key": "test.metadata.json",
+                    "LastModified": "2023-01-01T00:00:00+00:00",
+                }
+            ]
+        }
+    ]
+
+    mock_client.get_object.return_value = {"Body": mock_body}
 
     # First call should query IDrive e2
     await agent.async_list_backups()
-    assert mock_client.list_objects_v2.call_count == 1
+    assert mock_client.get_paginator.call_count == 1
     assert mock_client.get_object.call_count == 1
 
     # Second call should use cache
     await agent.async_list_backups()
-    assert mock_client.list_objects_v2.call_count == 1
+    assert mock_client.get_paginator.call_count == 1
     assert mock_client.get_object.call_count == 1
 
     # Set cache to expire
     agent._cache_expiration = time() - 1
 
     # Third call should query IDrive e2 again
     await agent.async_list_backups()
-    assert mock_client.list_objects_v2.call_count == 2
+    assert mock_client.get_paginator.call_count == 2
     assert mock_client.get_object.call_count == 2
 
 
@@ -526,3 +539,88 @@ async def test_listeners_get_cleaned_up(hass: HomeAssistant) -> None:
     remove_listener()
 
     assert DATA_BACKUP_AGENT_LISTENERS not in hass.data
+
+
+async def test_list_backups_with_pagination(
+    hass: HomeAssistant,
+    mock_config_entry: MockConfigEntry,
+) -> None:
+    """Test listing backups when paginating through multiple pages."""
+    # Create agent
+    agent = IDriveE2BackupAgent(hass, mock_config_entry)
+
+    # Create two different backups
+    backup1 = AgentBackup(
+        backup_id="backup1",
+        date="2023-01-01T00:00:00+00:00",
+        addons=[],
+        database_included=False,
+        extra_metadata={},
+        folders=[],
+        homeassistant_included=False,
+        homeassistant_version=None,
+        name="Backup 1",
+        protected=False,
+        size=0,
+    )
+    backup2 = AgentBackup(
+        backup_id="backup2",
+        date="2023-01-02T00:00:00+00:00",
+        addons=[],
+        database_included=False,
+        extra_metadata={},
+        folders=[],
+        homeassistant_included=False,
+        homeassistant_version=None,
+        name="Backup 2",
+        protected=False,
+        size=0,
+    )
+
+    # Setup two pages of results
+    page1 = {
+        "Contents": [
+            {
+                "Key": "backup1.metadata.json",
+                "LastModified": "2023-01-01T00:00:00+00:00",
+            },
+            {"Key": "backup1.tar", "LastModified": "2023-01-01T00:00:00+00:00"},
+        ]
+    }
+    page2 = {
+        "Contents": [
+            {
+                "Key": "backup2.metadata.json",
+                "LastModified": "2023-01-02T00:00:00+00:00",
+            },
+            {"Key": "backup2.tar", "LastModified": "2023-01-02T00:00:00+00:00"},
+        ]
+    }
+
+    # Setup mock client
+    mock_client = mock_config_entry.runtime_data
+    mock_client.get_paginator.return_value.paginate.return_value.__aiter__.return_value = [
+        page1,
+        page2,
+    ]
+
+    # Mock get_object responses based on the key
+    async def mock_get_object(**kwargs):
+        """Mock get_object with different responses based on the key."""
+        key = kwargs.get("Key", "")
+        if "backup1" in key:
+            mock_body = AsyncMock()
+            mock_body.read.return_value = json.dumps(backup1.as_dict()).encode()
+            return {"Body": mock_body}
+        # backup2
+        mock_body = AsyncMock()
+        mock_body.read.return_value = json.dumps(backup2.as_dict()).encode()
+        return {"Body": mock_body}
+
+    mock_client.get_object.side_effect = mock_get_object
+
+    # List backups and verify we got both
+    backups = await agent.async_list_backups()
+    assert len(backups) == 2
+    backup_ids = {backup.backup_id for backup in backups}
+    assert backup_ids == {"backup1", "backup2"}