From add184f47c18af9adaec5c0b462aacd023f4c4d8 Mon Sep 17 00:00:00 2001 From: Josef Zweck Date: Fri, 31 Jan 2025 08:23:49 +0000 Subject: [PATCH 1/2] Retry backup uploads in onedrive --- homeassistant/components/onedrive/backup.py | 31 +++-- tests/components/onedrive/conftest.py | 7 ++ tests/components/onedrive/test_backup.py | 121 +++++++++++++++++++- 3 files changed, 151 insertions(+), 8 deletions(-) diff --git a/homeassistant/components/onedrive/backup.py b/homeassistant/components/onedrive/backup.py index 94d60bc6398125..17d10e90f977ba 100644 --- a/homeassistant/components/onedrive/backup.py +++ b/homeassistant/components/onedrive/backup.py @@ -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 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,25 @@ 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, TimeoutException) as err: + if ( + isinstance(err, APIError) + and err.response_status_code + and err.response_status_code < 500 + ): # raise on 4xx errors + raise + if retries < MAX_RETRIES: + await asyncio.sleep(2**retries) + retries += 1 + continue + raise + retries = 0 start += UPLOAD_CHUNK_SIZE uploaded_chunks += 1 buffer_size -= UPLOAD_CHUNK_SIZE diff --git a/tests/components/onedrive/conftest.py b/tests/components/onedrive/conftest.py index 65142217017431..649966a7828978 100644 --- a/tests/components/onedrive/conftest.py +++ b/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 diff --git a/tests/components/onedrive/test_backup.py b/tests/components/onedrive/test_backup.py index 3492202d3feb08..7018abf0721d88 100644 --- a/tests/components/onedrive/test_backup.py +++ b/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,123 @@ async def test_broken_upload_session( assert "Failed to start backup upload" in caplog.text +async def test_agents_upload_errors_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=500), + 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 + + +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, +) -> 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=500) + + 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 "Backup operation failed" in caplog.text + + async def test_agents_download( hass_client: ClientSessionGenerator, mock_drive_items: MagicMock, @@ -282,7 +401,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( From 818e6c8914fb9219d2bf34da63ea819e7329a58e Mon Sep 17 00:00:00 2001 From: Josef Zweck Date: Fri, 31 Jan 2025 08:28:56 +0000 Subject: [PATCH 2/2] no exponential backup on timeout --- homeassistant/components/onedrive/backup.py | 13 +++++++----- tests/components/onedrive/test_backup.py | 23 ++++++++++++++++++--- 2 files changed, 28 insertions(+), 8 deletions(-) diff --git a/homeassistant/components/onedrive/backup.py b/homeassistant/components/onedrive/backup.py index 17d10e90f977ba..7f4bd5a07383cc 100644 --- a/homeassistant/components/onedrive/backup.py +++ b/homeassistant/components/onedrive/backup.py @@ -288,18 +288,21 @@ async def async_upload( start + UPLOAD_CHUNK_SIZE - 1, chunk_data[slice_start : slice_start + UPLOAD_CHUNK_SIZE], ) - except (APIError, TimeoutException) as err: + except APIError as err: if ( - isinstance(err, APIError) - and err.response_status_code - and err.response_status_code < 500 - ): # raise on 4xx errors + 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 diff --git a/tests/components/onedrive/test_backup.py b/tests/components/onedrive/test_backup.py index 7018abf0721d88..162ecb7d92ae74 100644 --- a/tests/components/onedrive/test_backup.py +++ b/tests/components/onedrive/test_backup.py @@ -257,19 +257,27 @@ 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 = [ - APIError(response_status_code=500), + side_effect, LargeFileUploadSession(next_expected_ranges=["2-"]), LargeFileUploadSession(next_expected_ranges=["2-"]), ] @@ -336,18 +344,27 @@ async def test_agents_upload_4xx_errors_not_retried( 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 = APIError(response_status_code=500) + mock_adapter.send_async.side_effect = side_effect with ( patch( @@ -371,7 +388,7 @@ async def test_agents_upload_fails_after_max_retries( 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 "Backup operation failed" in caplog.text + assert error in caplog.text async def test_agents_download(