Retry backup uploads in onedrive (#136980)

* Retry backup uploads in onedrive

* no exponential backup on timeout
This commit is contained in:
Josef Zweck 2025-01-31 11:23:33 +01:00 committed by Paulus Schoutsen
parent 26ae498974
commit 0272d37e88
3 changed files with 171 additions and 8 deletions

View File

@ -2,6 +2,7 @@
from __future__ import annotations
import asyncio
from collections.abc import AsyncIterator, Callable, Coroutine
from functools import wraps
import html
@ -9,7 +10,7 @@ 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 @@ from .const import DATA_BACKUP_AGENT_LISTENERS, DOMAIN
_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 @@ def handle_backup_errors[_R, **P](
)
_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 @@ class OneDriveBackupAgent(BackupAgent):
start = 0
buffer: list[bytes] = []
buffer_size = 0
retries = 0
async for chunk in stream:
buffer.append(chunk)
@ -279,11 +282,28 @@ class OneDriveBackupAgent(BackupAgent):
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

View File

@ -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

View File

@ -8,8 +8,10 @@ from io import StringIO
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(