Retry backup uploads in onedrive

homeassistant/components/onedrive/backup.py

@@ -2,14 +2,15 @@
 
 from __future__ import annotations
 
+import asyncio
 from collections.abc import AsyncIterator, Callable, Coroutine
 from functools import wraps
 import html
 import json
 import logging
 from typing import Any, Concatenate, cast
 
-from httpx import Response
+from httpx import Response, TimeoutException
 from kiota_abstractions.api_error import APIError
 from kiota_abstractions.authentication import AnonymousAuthenticationProvider
 from kiota_abstractions.headers_collection import HeadersCollection
@@ -42,6 +43,7 @@
 
 _LOGGER = logging.getLogger(__name__)
 UPLOAD_CHUNK_SIZE = 16 * 320 * 1024  # 5.2MB
+MAX_RETRIES = 5
 
 
 async def async_get_backup_agents(
@@ -96,7 +98,7 @@ async def wrapper(
             )
             _LOGGER.debug("Full error: %s", err, exc_info=True)
             raise BackupAgentError("Backup operation failed") from err
-        except TimeoutError as err:
+        except TimeoutException as err:
             _LOGGER.error(
                 "Error during backup in %s: Timeout",
                 func.__name__,
@@ -268,6 +270,7 @@ async def async_upload(
         start = 0
         buffer: list[bytes] = []
         buffer_size = 0
+        retries = 0
 
         async for chunk in stream:
             buffer.append(chunk)
@@ -279,11 +282,28 @@ async def async_upload(
                     buffer_size > UPLOAD_CHUNK_SIZE
                 ):  # Loop in case the buffer is >= UPLOAD_CHUNK_SIZE * 2
                     slice_start = uploaded_chunks * UPLOAD_CHUNK_SIZE
-                    await async_upload(
-                        start,
-                        start + UPLOAD_CHUNK_SIZE - 1,
-                        chunk_data[slice_start : slice_start + UPLOAD_CHUNK_SIZE],
-                    )
+                    try:
+                        await async_upload(
+                            start,
+                            start + UPLOAD_CHUNK_SIZE - 1,
+                            chunk_data[slice_start : slice_start + UPLOAD_CHUNK_SIZE],
+                        )
+                    except APIError as err:
+                        if (
+                            err.response_status_code and err.response_status_code < 500
+                        ):  # no retry on 4xx errors
+                            raise
+                        if retries < MAX_RETRIES:
+                            await asyncio.sleep(2**retries)
+                            retries += 1
+                            continue
+                        raise
+                    except TimeoutException:
+                        if retries < MAX_RETRIES:
+                            retries += 1
+                            continue
+                        raise
+                    retries = 0
                     start += UPLOAD_CHUNK_SIZE
                     uploaded_chunks += 1
                     buffer_size -= UPLOAD_CHUNK_SIZE

tests/components/onedrive/conftest.py

@@ -176,3 +176,10 @@ def mock_instance_id() -> Generator[AsyncMock]:
         return_value="9f86d081884c7d659a2feaa0c55ad015a3bf4f1b2b0",
     ):
         yield
+
+
+@pytest.fixture(autouse=True)
+def mock_asyncio_sleep() -> Generator[AsyncMock]:
+    """Mock asyncio.sleep."""
+    with patch("homeassistant.components.onedrive.backup.asyncio.sleep", AsyncMock()):
+        yield

tests/components/onedrive/test_backup.py

@@ -8,8 +8,10 @@
 from json import dumps
 from unittest.mock import Mock, patch
 
+from httpx import TimeoutException
 from kiota_abstractions.api_error import APIError
 from msgraph.generated.models.drive_item import DriveItem
+from msgraph_core.models import LargeFileUploadSession
 import pytest
 
 from homeassistant.components.backup import DOMAIN as BACKUP_DOMAIN, AgentBackup
@@ -255,6 +257,140 @@ async def test_broken_upload_session(
     assert "Failed to start backup upload" in caplog.text
 
 
+@pytest.mark.parametrize(
+    "side_effect",
+    [
+        APIError(response_status_code=500),
+        TimeoutException("Timeout"),
+    ],
+)
+async def test_agents_upload_errors_retried(
+    hass_client: ClientSessionGenerator,
+    caplog: pytest.LogCaptureFixture,
+    mock_drive_items: MagicMock,
+    mock_config_entry: MockConfigEntry,
+    mock_adapter: MagicMock,
+    side_effect: Exception,
+) -> None:
+    """Test agent upload backup."""
+    client = await hass_client()
+    test_backup = AgentBackup.from_dict(BACKUP_METADATA)
+
+    mock_adapter.send_async.side_effect = [
+        side_effect,
+        LargeFileUploadSession(next_expected_ranges=["2-"]),
+        LargeFileUploadSession(next_expected_ranges=["2-"]),
+    ]
+
+    with (
+        patch(
+            "homeassistant.components.backup.manager.BackupManager.async_get_backup",
+        ) as fetch_backup,
+        patch(
+            "homeassistant.components.backup.manager.read_backup",
+            return_value=test_backup,
+        ),
+        patch("pathlib.Path.open") as mocked_open,
+        patch("homeassistant.components.onedrive.backup.UPLOAD_CHUNK_SIZE", 3),
+    ):
+        mocked_open.return_value.read = Mock(side_effect=[b"test", b""])
+        fetch_backup.return_value = test_backup
+        resp = await client.post(
+            f"/api/backup/upload?agent_id={DOMAIN}.{mock_config_entry.unique_id}",
+            data={"file": StringIO("test")},
+        )
+
+    assert resp.status == 201
+    assert mock_adapter.send_async.call_count == 3
+    assert f"Uploading backup {test_backup.backup_id}" in caplog.text
+    mock_drive_items.patch.assert_called_once()
+
+
+async def test_agents_upload_4xx_errors_not_retried(
+    hass_client: ClientSessionGenerator,
+    caplog: pytest.LogCaptureFixture,
+    mock_drive_items: MagicMock,
+    mock_config_entry: MockConfigEntry,
+    mock_adapter: MagicMock,
+) -> None:
+    """Test agent upload backup."""
+    client = await hass_client()
+    test_backup = AgentBackup.from_dict(BACKUP_METADATA)
+
+    mock_adapter.send_async.side_effect = APIError(response_status_code=404)
+
+    with (
+        patch(
+            "homeassistant.components.backup.manager.BackupManager.async_get_backup",
+        ) as fetch_backup,
+        patch(
+            "homeassistant.components.backup.manager.read_backup",
+            return_value=test_backup,
+        ),
+        patch("pathlib.Path.open") as mocked_open,
+        patch("homeassistant.components.onedrive.backup.UPLOAD_CHUNK_SIZE", 3),
+    ):
+        mocked_open.return_value.read = Mock(side_effect=[b"test", b""])
+        fetch_backup.return_value = test_backup
+        resp = await client.post(
+            f"/api/backup/upload?agent_id={DOMAIN}.{mock_config_entry.unique_id}",
+            data={"file": StringIO("test")},
+        )
+
+    assert resp.status == 201
+    assert mock_adapter.send_async.call_count == 1
+    assert f"Uploading backup {test_backup.backup_id}" in caplog.text
+    assert mock_drive_items.patch.call_count == 0
+    assert "Backup operation failed" in caplog.text
+
+
+@pytest.mark.parametrize(
+    ("side_effect", "error"),
+    [
+        (APIError(response_status_code=500), "Backup operation failed"),
+        (TimeoutException("Timeout"), "Backup operation timed out"),
+    ],
+)
+async def test_agents_upload_fails_after_max_retries(
+    hass_client: ClientSessionGenerator,
+    caplog: pytest.LogCaptureFixture,
+    mock_drive_items: MagicMock,
+    mock_config_entry: MockConfigEntry,
+    mock_adapter: MagicMock,
+    side_effect: Exception,
+    error: str,
+) -> None:
+    """Test agent upload backup."""
+    client = await hass_client()
+    test_backup = AgentBackup.from_dict(BACKUP_METADATA)
+
+    mock_adapter.send_async.side_effect = side_effect
+
+    with (
+        patch(
+            "homeassistant.components.backup.manager.BackupManager.async_get_backup",
+        ) as fetch_backup,
+        patch(
+            "homeassistant.components.backup.manager.read_backup",
+            return_value=test_backup,
+        ),
+        patch("pathlib.Path.open") as mocked_open,
+        patch("homeassistant.components.onedrive.backup.UPLOAD_CHUNK_SIZE", 3),
+    ):
+        mocked_open.return_value.read = Mock(side_effect=[b"test", b""])
+        fetch_backup.return_value = test_backup
+        resp = await client.post(
+            f"/api/backup/upload?agent_id={DOMAIN}.{mock_config_entry.unique_id}",
+            data={"file": StringIO("test")},
+        )
+
+    assert resp.status == 201
+    assert mock_adapter.send_async.call_count == 6
+    assert f"Uploading backup {test_backup.backup_id}" in caplog.text
+    assert mock_drive_items.patch.call_count == 0
+    assert error in caplog.text
+
+
 async def test_agents_download(
     hass_client: ClientSessionGenerator,
     mock_drive_items: MagicMock,
@@ -282,7 +418,7 @@ async def test_agents_download(
             APIError(response_status_code=500),
             "Backup operation failed",
         ),
-        (TimeoutError(), "Backup operation timed out"),
+        (TimeoutException("Timeout"), "Backup operation timed out"),
     ],
 )
 async def test_delete_error(