diff --git a/homeassistant/components/aws_s3/backup.py b/homeassistant/components/aws_s3/backup.py index 784e267edab..95b9ae671b4 100644 --- a/homeassistant/components/aws_s3/backup.py +++ b/homeassistant/components/aws_s3/backup.py @@ -14,6 +14,7 @@ from homeassistant.components.backup import ( BackupAgent, BackupAgentError, BackupNotFound, + OnProgressCallback, suggested_filename, ) from homeassistant.core import HomeAssistant, callback @@ -132,6 +133,7 @@ class S3BackupAgent(BackupAgent): *, open_stream: Callable[[], Coroutine[Any, Any, AsyncIterator[bytes]]], backup: AgentBackup, + on_progress: OnProgressCallback, **kwargs: Any, ) -> None: """Upload a backup. diff --git a/homeassistant/components/azure_storage/backup.py b/homeassistant/components/azure_storage/backup.py index 54fd069a11f..5a684bfcc77 100644 --- a/homeassistant/components/azure_storage/backup.py +++ b/homeassistant/components/azure_storage/backup.py @@ -16,6 +16,7 @@ from homeassistant.components.backup import ( BackupAgent, BackupAgentError, BackupNotFound, + OnProgressCallback, suggested_filename, ) from homeassistant.core import HomeAssistant, callback @@ -129,6 +130,7 @@ class AzureStorageBackupAgent(BackupAgent): *, open_stream: Callable[[], Coroutine[Any, Any, AsyncIterator[bytes]]], backup: AgentBackup, + on_progress: OnProgressCallback, **kwargs: Any, ) -> None: """Upload a backup.""" diff --git a/homeassistant/components/backblaze_b2/backup.py b/homeassistant/components/backblaze_b2/backup.py index 9e795434c25..ec92a41a5dc 100644 --- a/homeassistant/components/backblaze_b2/backup.py +++ b/homeassistant/components/backblaze_b2/backup.py @@ -17,6 +17,7 @@ from homeassistant.components.backup import ( BackupAgent, BackupAgentError, BackupNotFound, + OnProgressCallback, suggested_filename, ) from homeassistant.core import HomeAssistant, callback @@ -230,6 +231,7 @@ class BackblazeBackupAgent(BackupAgent): *, open_stream: Callable[[], Coroutine[Any, Any, AsyncIterator[bytes]]], backup: AgentBackup, + on_progress: OnProgressCallback, **kwargs: Any, ) -> None: """Upload a backup to Backblaze B2. diff --git a/homeassistant/components/backup/__init__.py b/homeassistant/components/backup/__init__.py index f3289d6e744..6ed4f1ac2d3 100644 --- a/homeassistant/components/backup/__init__.py +++ b/homeassistant/components/backup/__init__.py @@ -17,6 +17,7 @@ from .agent import ( BackupAgentError, BackupAgentPlatformProtocol, LocalBackupAgent, + OnProgressCallback, ) from .config import BackupConfig, CreateBackupParametersDict from .const import DATA_MANAGER, DOMAIN @@ -41,6 +42,7 @@ from .manager import ( RestoreBackupEvent, RestoreBackupStage, RestoreBackupState, + UploadBackupEvent, WrittenBackup, ) from .models import AddonInfo, AgentBackup, BackupNotFound, Folder @@ -72,9 +74,11 @@ __all__ = [ "LocalBackupAgent", "ManagerBackup", "NewBackup", + "OnProgressCallback", "RestoreBackupEvent", "RestoreBackupStage", "RestoreBackupState", + "UploadBackupEvent", "WrittenBackup", "async_get_manager", "suggested_filename", diff --git a/homeassistant/components/backup/agent.py b/homeassistant/components/backup/agent.py index 8093ac88338..afb4cbf1d18 100644 --- a/homeassistant/components/backup/agent.py +++ b/homeassistant/components/backup/agent.py @@ -14,6 +14,13 @@ from homeassistant.core import HomeAssistant, callback from .models import AgentBackup, BackupAgentError +class OnProgressCallback(Protocol): + """Protocol for on_progress callback.""" + + def __call__(self, *, bytes_uploaded: int, **kwargs: Any) -> None: + """Report upload progress.""" + + class BackupAgentUnreachableError(BackupAgentError): """Raised when the agent can't reach its API.""" @@ -53,12 +60,14 @@ class BackupAgent(abc.ABC): *, open_stream: Callable[[], Coroutine[Any, Any, AsyncIterator[bytes]]], backup: AgentBackup, + on_progress: OnProgressCallback, **kwargs: Any, ) -> None: """Upload a backup. :param open_stream: A function returning an async iterator that yields bytes. :param backup: Metadata about the backup that should be uploaded. + :param on_progress: A callback to report the number of uploaded bytes. """ @abc.abstractmethod diff --git a/homeassistant/components/backup/backup.py b/homeassistant/components/backup/backup.py index de2cfecb1a5..3396c7e103f 100644 --- a/homeassistant/components/backup/backup.py +++ b/homeassistant/components/backup/backup.py @@ -11,7 +11,7 @@ from typing import Any from homeassistant.core import HomeAssistant from homeassistant.helpers.hassio import is_hassio -from .agent import BackupAgent, LocalBackupAgent +from .agent import BackupAgent, LocalBackupAgent, OnProgressCallback from .const import DOMAIN, LOGGER from .models import AgentBackup, BackupNotFound from .util import read_backup, suggested_filename @@ -73,6 +73,7 @@ class CoreLocalBackupAgent(LocalBackupAgent): *, open_stream: Callable[[], Coroutine[Any, Any, AsyncIterator[bytes]]], backup: AgentBackup, + on_progress: OnProgressCallback, **kwargs: Any, ) -> None: """Upload a backup.""" diff --git a/homeassistant/components/backup/manager.py b/homeassistant/components/backup/manager.py index 909225f5bde..fbd73a31923 100644 --- a/homeassistant/components/backup/manager.py +++ b/homeassistant/components/backup/manager.py @@ -252,6 +252,15 @@ class BlockedEvent(ManagerStateEvent): manager_state: BackupManagerState = BackupManagerState.BLOCKED +@dataclass(frozen=True, kw_only=True, slots=True) +class UploadBackupEvent(ManagerStateEvent): + """Backup agent upload progress event.""" + + agent_id: str + uploaded_bytes: int + total_bytes: int + + class BackupPlatformProtocol(Protocol): """Define the format that backup platforms can have.""" @@ -579,9 +588,24 @@ class BackupManager: _backup = replace( backup, protected=should_encrypt, size=streamer.size() ) - await self.backup_agents[agent_id].async_upload_backup( + agent = self.backup_agents[agent_id] + + @callback + def on_upload_progress(*, bytes_uploaded: int, **kwargs: Any) -> None: + """Handle upload progress.""" + self.async_on_backup_event( + UploadBackupEvent( + manager_state=self.state, + agent_id=agent_id, + uploaded_bytes=bytes_uploaded, + total_bytes=_backup.size, + ) + ) + + await agent.async_upload_backup( open_stream=open_stream_func, backup=_backup, + on_progress=on_upload_progress, ) if streamer: await streamer.wait() @@ -1374,9 +1398,10 @@ class BackupManager: """Forward event to subscribers.""" if (current_state := self.state) != (new_state := event.manager_state): LOGGER.debug("Backup state: %s -> %s", current_state, new_state) - self.last_event = event - if not isinstance(event, (BlockedEvent, IdleEvent)): - self.last_action_event = event + if not isinstance(event, UploadBackupEvent): + self.last_event = event + if not isinstance(event, (BlockedEvent, IdleEvent)): + self.last_action_event = event for subscription in self._backup_event_subscriptions: subscription(event) diff --git a/homeassistant/components/cloud/backup.py b/homeassistant/components/cloud/backup.py index bca65a68abd..180c14ef111 100644 --- a/homeassistant/components/cloud/backup.py +++ b/homeassistant/components/cloud/backup.py @@ -18,6 +18,7 @@ from homeassistant.components.backup import ( BackupAgent, BackupAgentError, BackupNotFound, + OnProgressCallback, ) from homeassistant.core import HomeAssistant, callback from homeassistant.helpers.aiohttp_client import ChunkAsyncStreamIterator @@ -106,6 +107,7 @@ class CloudBackupAgent(BackupAgent): *, open_stream: Callable[[], Coroutine[Any, Any, AsyncIterator[bytes]]], backup: AgentBackup, + on_progress: OnProgressCallback, **kwargs: Any, ) -> None: """Upload a backup. diff --git a/homeassistant/components/cloudflare_r2/backup.py b/homeassistant/components/cloudflare_r2/backup.py index cef9294182e..4fc8199a4b3 100644 --- a/homeassistant/components/cloudflare_r2/backup.py +++ b/homeassistant/components/cloudflare_r2/backup.py @@ -14,6 +14,7 @@ from homeassistant.components.backup import ( BackupAgent, BackupAgentError, BackupNotFound, + OnProgressCallback, suggested_filename, ) from homeassistant.core import HomeAssistant, callback @@ -129,6 +130,7 @@ class R2BackupAgent(BackupAgent): *, open_stream: Callable[[], Coroutine[Any, Any, AsyncIterator[bytes]]], backup: AgentBackup, + on_progress: OnProgressCallback, **kwargs: Any, ) -> None: """Upload a backup. diff --git a/homeassistant/components/google_drive/backup.py b/homeassistant/components/google_drive/backup.py index bc306fe61d7..e6967d95eaf 100644 --- a/homeassistant/components/google_drive/backup.py +++ b/homeassistant/components/google_drive/backup.py @@ -13,6 +13,7 @@ from homeassistant.components.backup import ( BackupAgent, BackupAgentError, BackupNotFound, + OnProgressCallback, ) from homeassistant.core import HomeAssistant, callback from homeassistant.exceptions import HomeAssistantError @@ -75,6 +76,7 @@ class GoogleDriveBackupAgent(BackupAgent): *, open_stream: Callable[[], Coroutine[Any, Any, AsyncIterator[bytes]]], backup: AgentBackup, + on_progress: OnProgressCallback, **kwargs: Any, ) -> None: """Upload a backup. diff --git a/homeassistant/components/hassio/backup.py b/homeassistant/components/hassio/backup.py index 1e9a14be1f2..aeafe7d4d96 100644 --- a/homeassistant/components/hassio/backup.py +++ b/homeassistant/components/hassio/backup.py @@ -44,6 +44,7 @@ from homeassistant.components.backup import ( IncorrectPasswordError, ManagerBackup, NewBackup, + OnProgressCallback, RestoreBackupEvent, RestoreBackupStage, RestoreBackupState, @@ -183,6 +184,7 @@ class SupervisorBackupAgent(BackupAgent): *, open_stream: Callable[[], Coroutine[Any, Any, AsyncIterator[bytes]]], backup: AgentBackup, + on_progress: OnProgressCallback, **kwargs: Any, ) -> None: """Upload a backup. diff --git a/homeassistant/components/idrive_e2/backup.py b/homeassistant/components/idrive_e2/backup.py index 4df337fa27b..2fcdcf73ecc 100644 --- a/homeassistant/components/idrive_e2/backup.py +++ b/homeassistant/components/idrive_e2/backup.py @@ -15,6 +15,7 @@ from homeassistant.components.backup import ( BackupAgent, BackupAgentError, BackupNotFound, + OnProgressCallback, suggested_filename, ) from homeassistant.core import HomeAssistant, callback @@ -127,6 +128,7 @@ class IDriveE2BackupAgent(BackupAgent): *, open_stream: Callable[[], Coroutine[Any, Any, AsyncIterator[bytes]]], backup: AgentBackup, + on_progress: OnProgressCallback, **kwargs: Any, ) -> None: """Upload a backup. diff --git a/homeassistant/components/kitchen_sink/backup.py b/homeassistant/components/kitchen_sink/backup.py index 46b204845ad..1ff9cc5e05d 100644 --- a/homeassistant/components/kitchen_sink/backup.py +++ b/homeassistant/components/kitchen_sink/backup.py @@ -13,6 +13,7 @@ from homeassistant.components.backup import ( BackupAgent, BackupNotFound, Folder, + OnProgressCallback, ) from homeassistant.core import HomeAssistant, callback @@ -91,6 +92,7 @@ class KitchenSinkBackupAgent(BackupAgent): *, open_stream: Callable[[], Coroutine[Any, Any, AsyncIterator[bytes]]], backup: AgentBackup, + on_progress: OnProgressCallback, **kwargs: Any, ) -> None: """Upload a backup.""" diff --git a/homeassistant/components/onedrive/backup.py b/homeassistant/components/onedrive/backup.py index a76d6df820a..5dd7038f211 100644 --- a/homeassistant/components/onedrive/backup.py +++ b/homeassistant/components/onedrive/backup.py @@ -22,6 +22,7 @@ from homeassistant.components.backup import ( BackupAgent, BackupAgentError, BackupNotFound, + OnProgressCallback, suggested_filename, ) from homeassistant.core import HomeAssistant, callback @@ -145,6 +146,7 @@ class OneDriveBackupAgent(BackupAgent): *, open_stream: Callable[[], Coroutine[Any, Any, AsyncIterator[bytes]]], backup: AgentBackup, + on_progress: OnProgressCallback, **kwargs: Any, ) -> None: """Upload a backup.""" diff --git a/homeassistant/components/onedrive_for_business/backup.py b/homeassistant/components/onedrive_for_business/backup.py index 52ce8af8941..bec7dfd8c3e 100644 --- a/homeassistant/components/onedrive_for_business/backup.py +++ b/homeassistant/components/onedrive_for_business/backup.py @@ -22,6 +22,7 @@ from homeassistant.components.backup import ( BackupAgent, BackupAgentError, BackupNotFound, + OnProgressCallback, suggested_filename, ) from homeassistant.core import HomeAssistant, callback @@ -145,6 +146,7 @@ class OneDriveBackupAgent(BackupAgent): *, open_stream: Callable[[], Coroutine[Any, Any, AsyncIterator[bytes]]], backup: AgentBackup, + on_progress: OnProgressCallback, **kwargs: Any, ) -> None: """Upload a backup.""" diff --git a/homeassistant/components/sftp_storage/backup.py b/homeassistant/components/sftp_storage/backup.py index 4859f2d2f2a..2367d022a44 100644 --- a/homeassistant/components/sftp_storage/backup.py +++ b/homeassistant/components/sftp_storage/backup.py @@ -12,6 +12,7 @@ from homeassistant.components.backup import ( BackupAgent, BackupAgentError, BackupNotFound, + OnProgressCallback, ) from homeassistant.core import HomeAssistant, callback @@ -85,6 +86,7 @@ class SFTPBackupAgent(BackupAgent): *, open_stream: Callable[[], Coroutine[Any, Any, AsyncIterator[bytes]]], backup: AgentBackup, + on_progress: OnProgressCallback, **kwargs: Any, ) -> None: """Upload a backup.""" diff --git a/homeassistant/components/synology_dsm/backup.py b/homeassistant/components/synology_dsm/backup.py index b3279db1cac..3933a3f2fc2 100644 --- a/homeassistant/components/synology_dsm/backup.py +++ b/homeassistant/components/synology_dsm/backup.py @@ -15,6 +15,7 @@ from homeassistant.components.backup import ( BackupAgent, BackupAgentError, BackupNotFound, + OnProgressCallback, suggested_filename, ) from homeassistant.core import HomeAssistant, callback @@ -155,6 +156,7 @@ class SynologyDSMBackupAgent(BackupAgent): *, open_stream: Callable[[], Coroutine[Any, Any, AsyncIterator[bytes]]], backup: AgentBackup, + on_progress: OnProgressCallback, **kwargs: Any, ) -> None: """Upload a backup. diff --git a/homeassistant/components/webdav/backup.py b/homeassistant/components/webdav/backup.py index 5b27e7be29b..6f856d5de5e 100644 --- a/homeassistant/components/webdav/backup.py +++ b/homeassistant/components/webdav/backup.py @@ -17,6 +17,7 @@ from homeassistant.components.backup import ( BackupAgent, BackupAgentError, BackupNotFound, + OnProgressCallback, suggested_filename, ) from homeassistant.core import HomeAssistant, callback @@ -140,6 +141,7 @@ class WebDavBackupAgent(BackupAgent): *, open_stream: Callable[[], Coroutine[Any, Any, AsyncIterator[bytes]]], backup: AgentBackup, + on_progress: OnProgressCallback, **kwargs: Any, ) -> None: """Upload a backup. diff --git a/tests/components/backup/common.py b/tests/components/backup/common.py index d9533d2764d..55f15f628d7 100644 --- a/tests/components/backup/common.py +++ b/tests/components/backup/common.py @@ -173,6 +173,7 @@ async def setup_backup_integration( side_effect=RuntimeError("Local agent does not open stream") ), backup=backup, + on_progress=lambda *, on_progress, **_: None, ) local_agent._loaded_backups = True diff --git a/tests/components/backup/test_manager.py b/tests/components/backup/test_manager.py index 73cb98d7fd3..b8b69df749e 100644 --- a/tests/components/backup/test_manager.py +++ b/tests/components/backup/test_manager.py @@ -3709,3 +3709,77 @@ async def test_manager_not_blocked_after_restore( "next_automatic_backup_additional": False, "state": "idle", } + + +async def test_upload_progress_event( + hass: HomeAssistant, + hass_ws_client: WebSocketGenerator, + generate_backup_id: MagicMock, +) -> None: + """Test that upload progress events are fired when an agent reports progress.""" + agent_ids = [LOCAL_AGENT_ID, "test.remote"] + mock_agents = await setup_backup_integration(hass, remote_agents=["test.remote"]) + + remote_agent = mock_agents["test.remote"] + original_side_effect = remote_agent.async_upload_backup.side_effect + + async def upload_with_progress(**kwargs: Any) -> None: + """Upload and report progress.""" + on_progress = kwargs["on_progress"] + on_progress(bytes_uploaded=500) + on_progress(bytes_uploaded=1000) + await original_side_effect(**kwargs) + + remote_agent.async_upload_backup.side_effect = upload_with_progress + + ws_client = await hass_ws_client(hass) + + await ws_client.send_json_auto_id({"type": "backup/subscribe_events"}) + + result = await ws_client.receive_json() + assert result["event"] == {"manager_state": BackupManagerState.IDLE} + + result = await ws_client.receive_json() + assert result["success"] is True + + with patch("pathlib.Path.open", mock_open(read_data=b"test")): + await ws_client.send_json_auto_id( + {"type": "backup/generate", "agent_ids": agent_ids} + ) + result = await ws_client.receive_json() + assert result["event"]["manager_state"] == BackupManagerState.CREATE_BACKUP + + result = await ws_client.receive_json() + assert result["success"] is True + + await hass.async_block_till_done() + + # Consume intermediate stage events (home_assistant, upload_to_agents) + result = await ws_client.receive_json() + assert result["event"]["stage"] == CreateBackupStage.HOME_ASSISTANT + + result = await ws_client.receive_json() + assert result["event"]["stage"] == CreateBackupStage.UPLOAD_TO_AGENTS + + # Upload progress events for the remote agent + result = await ws_client.receive_json() + assert result["event"] == { + "manager_state": BackupManagerState.CREATE_BACKUP, + "agent_id": "test.remote", + "uploaded_bytes": 500, + "total_bytes": ANY, + } + + result = await ws_client.receive_json() + assert result["event"] == { + "manager_state": BackupManagerState.CREATE_BACKUP, + "agent_id": "test.remote", + "uploaded_bytes": 1000, + "total_bytes": ANY, + } + + result = await ws_client.receive_json() + assert result["event"]["state"] == CreateBackupState.COMPLETED + + result = await ws_client.receive_json() + assert result["event"] == {"manager_state": BackupManagerState.IDLE}