diff --git a/homeassistant/components/google_drive/api.py b/homeassistant/components/google_drive/api.py index f6dfaa1b76bc..10d51a5fa292 100644 --- a/homeassistant/components/google_drive/api.py +++ b/homeassistant/components/google_drive/api.py @@ -34,6 +34,43 @@ class StorageQuotaData: usage_in_trash: int +def _invalid_metadata_reason(metadata: Any) -> str | None: + """Return why decoded metadata cannot be used, or None if it can. + + AgentBackup.from_dict does not enforce its annotations, so the types the + backup manager relies on are checked here, before it is built: the manager + uses backup_id as a dict key and calls extra_metadata.get() on every backup. + """ + if not isinstance(metadata, dict): + return "description is not a JSON object" + if not isinstance(metadata.get("backup_id"), str): + return "backup_id is not a string" + if not isinstance(metadata.get("extra_metadata"), dict): + return "extra_metadata is not a dictionary" + return None + + +def _parse_backup_metadata(file: dict[str, Any]) -> AgentBackup | None: + """Return the backup a Drive file describes, or None if it cannot be read. + + The metadata lives in the file description, which the user can edit or clear + from the Google Drive UI. One unreadable file should not hide the others. + """ + reason: object + try: + metadata = json.loads(file["description"]) + if (reason := _invalid_metadata_reason(metadata)) is None: + return AgentBackup.from_dict(metadata) + except (KeyError, TypeError, ValueError) as err: + reason = err + _LOGGER.warning( + "Ignoring backup file %s: its description is not valid backup metadata: %s", + file.get("id", "?"), + reason, + ) + return None + + class AsyncConfigEntryAuth(AbstractAuth): """Provide Google Drive authentication tied to an OAuth2 based config entry.""" @@ -195,42 +232,49 @@ class DriveClient: backup_metadata["name"], ) - async def async_list_backups(self) -> list[AgentBackup]: - """List backups.""" - query = " and ".join( + def _backup_query(self, *extra: str) -> str: + """Return a query matching the backups of this Home Assistant instance.""" + return " and ".join( [ "properties has { key='home_assistant' and value='backup' }", "properties has { key='instance_id'" f" and value='{self._ha_instance_id}' }}", "trashed=false", + *extra, ] ) + + async def async_list_backups(self) -> list[AgentBackup]: + """List backups.""" res = await self._api.list_files( - params={"q": query, "fields": "files(description)"} + params={"q": self._backup_query(), "fields": "files(id,description)"} ) - backups = [] - for file in res["files"]: - backup = AgentBackup.from_dict(json.loads(file["description"])) - backups.append(backup) - return backups + return [ + backup + for file in res["files"] + if (backup := _parse_backup_metadata(file)) is not None + ] async def async_get_size_of_all_backups(self) -> int: """Get size of all backups.""" - backups = await self.async_list_backups() - - return sum(backup.size for backup in backups) + # Ask Drive for the size of each file instead of adding up the sizes stored + # in the metadata, which would mean downloading and parsing every backup's + # description just to update a sensor. + res = await self._api.list_files( + params={"q": self._backup_query(), "fields": "files(size)"} + ) + return sum(int(file["size"]) for file in res["files"] if "size" in file) async def async_get_backup_file_id(self, backup_id: str) -> str | None: """Get file_id of backup if it exists.""" - query = " and ".join( - [ - "properties has { key='home_assistant' and value='backup' }", - "properties has { key='instance_id'" - f" and value='{self._ha_instance_id}' }}", - f"properties has {{ key='backup_id' and value='{backup_id}' }}", - ] + res = await self._api.list_files( + params={ + "q": self._backup_query( + f"properties has {{ key='backup_id' and value='{backup_id}' }}" + ), + "fields": "files(id)", + } ) - res = await self._api.list_files(params={"q": query, "fields": "files(id)"}) for file in res["files"]: return str(file["id"]) return None diff --git a/homeassistant/components/google_drive/backup.py b/homeassistant/components/google_drive/backup.py index 8f50d5208304..3be2972cf29f 100644 --- a/homeassistant/components/google_drive/backup.py +++ b/homeassistant/components/google_drive/backup.py @@ -85,16 +85,24 @@ class GoogleDriveBackupAgent(BackupAgent): :param backup: Metadata about the backup that should be uploaded. """ + bytes_uploaded = 0 + @wraps(open_stream) async def wrapped_open_stream() -> AsyncIterator[bytes]: stream = await open_stream() async def _progress_stream() -> AsyncIterator[bytes]: - bytes_uploaded = 0 + nonlocal bytes_uploaded + position = 0 async for chunk in stream: yield chunk - bytes_uploaded += len(chunk) - on_progress(bytes_uploaded=bytes_uploaded) + position += len(chunk) + # A retried upload reopens the stream from the beginning and + # skips whatever the server already received, so only report + # progress once it passes what was previously uploaded. + if position > bytes_uploaded: + bytes_uploaded = position + on_progress(bytes_uploaded=bytes_uploaded) return _progress_stream() diff --git a/tests/components/google_drive/snapshots/test_backup.ambr b/tests/components/google_drive/snapshots/test_backup.ambr index d783dad34d04..c9e31e3b53be 100644 --- a/tests/components/google_drive/snapshots/test_backup.ambr +++ b/tests/components/google_drive/snapshots/test_backup.ambr @@ -38,7 +38,7 @@ ), dict({ 'params': dict({ - 'fields': 'files(description)', + 'fields': 'files(size)', 'q': "properties has { key='home_assistant' and value='backup' } and properties has { key='instance_id' and value='0a123c' } and trashed=false", }), }), @@ -50,7 +50,7 @@ dict({ 'params': dict({ 'fields': 'files(id)', - 'q': "properties has { key='home_assistant' and value='backup' } and properties has { key='instance_id' and value='0a123c' } and properties has { key='backup_id' and value='test-backup' }", + 'q': "properties has { key='home_assistant' and value='backup' } and properties has { key='instance_id' and value='0a123c' } and trashed=false and properties has { key='backup_id' and value='test-backup' }", }), }), ), @@ -103,7 +103,7 @@ ), dict({ 'params': dict({ - 'fields': 'files(description)', + 'fields': 'files(size)', 'q': "properties has { key='home_assistant' and value='backup' } and properties has { key='instance_id' and value='0a123c' } and trashed=false", }), }), @@ -114,7 +114,7 @@ ), dict({ 'params': dict({ - 'fields': 'files(description)', + 'fields': 'files(id,description)', 'q': "properties has { key='home_assistant' and value='backup' } and properties has { key='instance_id' and value='0a123c' } and trashed=false", }), }), @@ -126,7 +126,7 @@ dict({ 'params': dict({ 'fields': 'files(id)', - 'q': "properties has { key='home_assistant' and value='backup' } and properties has { key='instance_id' and value='0a123c' } and properties has { key='backup_id' and value='test-backup' }", + 'q': "properties has { key='home_assistant' and value='backup' } and properties has { key='instance_id' and value='0a123c' } and trashed=false and properties has { key='backup_id' and value='test-backup' }", }), }), ), @@ -186,7 +186,7 @@ ), dict({ 'params': dict({ - 'fields': 'files(description)', + 'fields': 'files(size)', 'q': "properties has { key='home_assistant' and value='backup' } and properties has { key='instance_id' and value='0a123c' } and trashed=false", }), }), @@ -197,7 +197,7 @@ ), dict({ 'params': dict({ - 'fields': 'files(description)', + 'fields': 'files(id,description)', 'q': "properties has { key='home_assistant' and value='backup' } and properties has { key='instance_id' and value='0a123c' } and trashed=false", }), }), @@ -243,7 +243,7 @@ ), dict({ 'params': dict({ - 'fields': 'files(description)', + 'fields': 'files(size)', 'q': "properties has { key='home_assistant' and value='backup' } and properties has { key='instance_id' and value='0a123c' } and trashed=false", }), }), @@ -329,7 +329,7 @@ ), dict({ 'params': dict({ - 'fields': 'files(description)', + 'fields': 'files(size)', 'q': "properties has { key='home_assistant' and value='backup' } and properties has { key='instance_id' and value='0a123c' } and trashed=false", }), }), diff --git a/tests/components/google_drive/snapshots/test_diagnostic.ambr b/tests/components/google_drive/snapshots/test_diagnostic.ambr index d88e3da02bd0..eb3c732db4c7 100644 --- a/tests/components/google_drive/snapshots/test_diagnostic.ambr +++ b/tests/components/google_drive/snapshots/test_diagnostic.ambr @@ -41,7 +41,7 @@ }), }), 'coordinator_data': dict({ - 'all_backups_size': 104857600.0, + 'all_backups_size': 104857600, 'storage_quota': dict({ 'limit': 10737418240, 'usage': 5368709120, diff --git a/tests/components/google_drive/test_backup.py b/tests/components/google_drive/test_backup.py index 48e2b72878a3..85e6a83ae59f 100644 --- a/tests/components/google_drive/test_backup.py +++ b/tests/components/google_drive/test_backup.py @@ -73,6 +73,23 @@ async def consume_stream( pass +async def consume_stream_twice( + file_metadata: Any, + open_stream: Any, + *args: Any, + **kwargs: Any, +) -> None: + """Consume the stream twice, like a resumable upload that had to retry. + + A retried upload reopens the stream from the beginning and skips whatever + the server already received. + """ + for _ in range(2): + stream = await open_stream() + async for _ in stream: + pass + + @pytest.fixture(autouse=True) async def setup_integration( hass: HomeAssistant, @@ -143,6 +160,84 @@ async def test_agents_list_backups( assert [tuple(mock_call) for mock_call in mock_api.mock_calls] == snapshot +async def test_agents_list_backups_ignores_unreadable_metadata( + hass: HomeAssistant, + hass_ws_client: WebSocketGenerator, + mock_api: MagicMock, + caplog: pytest.LogCaptureFixture, +) -> None: + """Test that a backup whose description cannot be read is skipped. + + The description is editable from the Google Drive UI, so one unreadable + file must not hide the others. + """ + mock_api.list_files = AsyncMock( + return_value={ + "files": [ + {"id": "no description at all"}, + {"id": "not json", "description": "cleared by the user"}, + {"id": "not backup metadata", "description": '{"foo": "bar"}'}, + {"description": json.dumps(TEST_AGENT_BACKUP.as_dict())}, + ] + } + ) + + client = await hass_ws_client(hass) + await client.send_json_auto_id({"type": "backup/info"}) + response = await client.receive_json() + + assert response["success"] + assert response["result"]["agent_errors"] == {} + assert response["result"]["backups"] == [TEST_AGENT_BACKUP_RESULT] + assert "Ignoring backup file no description at all" in caplog.text + assert "Ignoring backup file not json" in caplog.text + assert "Ignoring backup file not backup metadata" in caplog.text + + +@pytest.mark.parametrize( + ("field", "value"), + [ + # The backup manager calls extra_metadata.get() on every backup. + ("extra_metadata", []), + # The backup manager uses backup_id as a dict key. + ("backup_id", ["not a string"]), + ], +) +async def test_agents_list_backups_ignores_wrong_typed_metadata( + hass: HomeAssistant, + hass_ws_client: WebSocketGenerator, + mock_api: MagicMock, + caplog: pytest.LogCaptureFixture, + field: str, + value: Any, +) -> None: + """Test that metadata which decodes but has the wrong types is skipped. + + AgentBackup.from_dict does not enforce its annotations, so such a backup is + only rejected once the backup manager uses it. + """ + wrong_types = TEST_AGENT_BACKUP.as_dict() + wrong_types[field] = value + mock_api.list_files = AsyncMock( + return_value={ + "files": [ + {"id": "wrong types", "description": json.dumps(wrong_types)}, + {"description": json.dumps(TEST_AGENT_BACKUP.as_dict())}, + ] + } + ) + + client = await hass_ws_client(hass) + await client.send_json_auto_id({"type": "backup/info"}) + response = await client.receive_json() + + assert response["success"] + assert response["result"]["agent_errors"] == {} + assert response["result"]["backups"] == [TEST_AGENT_BACKUP_RESULT] + assert "Ignoring backup file wrong types" in caplog.text + assert f"{field} is not a" in caplog.text + + async def test_agents_list_backups_fail( hass: HomeAssistant, hass_ws_client: WebSocketGenerator, @@ -399,6 +494,37 @@ async def test_agents_upload_progress( assert progress_calls == [6, 12] +async def test_agents_upload_progress_does_not_go_backwards_on_retry( + hass: HomeAssistant, + mock_api: MagicMock, +) -> None: + """Test agent upload progress is not reported twice when the upload retries.""" + mock_api.resumable_upload_file = AsyncMock(side_effect=consume_stream_twice) + + entries = hass.config_entries.async_entries(DOMAIN) + agent = GoogleDriveBackupAgent(entries[0]) + + progress_calls = [] + + def on_progress(*, bytes_uploaded: int, **kwargs: Any) -> None: + progress_calls.append(bytes_uploaded) + + async def open_stream() -> AsyncIterator[bytes]: + async def stream() -> AsyncIterator[bytes]: + yield b"chunk1" + yield b"chunk2" + + return stream() + + await agent.async_upload_backup( + open_stream=open_stream, + backup=TEST_AGENT_BACKUP, + on_progress=on_progress, + ) + + assert progress_calls == [6, 12] + + async def test_agents_upload_fail( hass: HomeAssistant, hass_client: ClientSessionGenerator, diff --git a/tests/components/google_drive/test_diagnostic.py b/tests/components/google_drive/test_diagnostic.py index 006df85b5152..e2d0c83f84dc 100644 --- a/tests/components/google_drive/test_diagnostic.py +++ b/tests/components/google_drive/test_diagnostic.py @@ -35,6 +35,7 @@ async def test_entry_diagnostics( "id": "HA folder ID", "name": "HA folder name", "description": json.dumps(mock_agent_backup.as_dict()), + "size": str(int(mock_agent_backup.size)), } ] } diff --git a/tests/components/google_drive/test_sensor.py b/tests/components/google_drive/test_sensor.py index 6e9d959cfbcd..dab2aa6f2edf 100644 --- a/tests/components/google_drive/test_sensor.py +++ b/tests/components/google_drive/test_sensor.py @@ -1,6 +1,5 @@ """Tests for the Google Drive sensor platform.""" -import json from unittest.mock import AsyncMock, MagicMock from freezegun.api import FrozenDateTimeFactory @@ -123,15 +122,7 @@ async def test_calculate_backups_size( assert state.state == "0.0" mock_api.list_files = AsyncMock( - return_value={ - "files": [ - { - "id": "HA folder ID", - "name": "HA folder name", - "description": json.dumps(mock_agent_backup.as_dict()), - } - ] - } + return_value={"files": [{"size": str(int(mock_agent_backup.size))}]} ) freezer.tick(SCAN_INTERVAL) async_fire_time_changed(hass) @@ -141,3 +132,26 @@ async def test_calculate_backups_size( state := hass.states.get("sensor.testuser_domain_com_total_size_of_backups") ) assert state.state == "100.0" + + +@pytest.mark.usefixtures("entity_registry_enabled_by_default") +async def test_calculate_backups_size_ignores_files_without_size( + hass: HomeAssistant, + mock_api: MagicMock, + config_entry: MockConfigEntry, + freezer: FrozenDateTimeFactory, +) -> None: + """Test that a file Google Drive reports no size for is skipped.""" + await setup_integration(hass, config_entry) + + mock_api.list_files = AsyncMock( + return_value={"files": [{"size": "1048576"}, {"id": "no size reported"}]} + ) + freezer.tick(SCAN_INTERVAL) + async_fire_time_changed(hass) + await hass.async_block_till_done() + + assert ( + state := hass.states.get("sensor.testuser_domain_com_total_size_of_backups") + ) + assert state.state == "1.0"