"""Test the AWS S3 backup platform.""" from collections.abc import AsyncGenerator from io import StringIO import json from time import time from unittest.mock import ANY, AsyncMock, Mock, call, patch from botocore.exceptions import ConnectTimeoutError import pytest from homeassistant.components.aws_s3.backup import ( MULTIPART_MIN_PART_SIZE_BYTES, BotoCoreError, S3BackupAgent, async_register_backup_agents_listener, suggested_filenames, ) from homeassistant.components.aws_s3.const import ( CONF_ENDPOINT_URL, DATA_BACKUP_AGENT_LISTENERS, DOMAIN, ) from homeassistant.components.backup import ( DATA_MANAGER, DOMAIN as BACKUP_DOMAIN, AgentBackup, UploadBackupEvent, ) from homeassistant.core import HomeAssistant from homeassistant.setup import async_setup_component from . import setup_integration from .const import USER_INPUT from tests.common import MockConfigEntry from tests.typing import ClientSessionGenerator, MagicMock, WebSocketGenerator @pytest.fixture(autouse=True) async def setup_backup_integration( hass: HomeAssistant, mock_config_entry: MockConfigEntry, ) -> AsyncGenerator[None]: """Set up S3 integration.""" with ( patch("homeassistant.components.backup.is_hassio", return_value=False), patch("homeassistant.components.backup.store.STORE_DELAY_SAVE", 0), ): assert await async_setup_component(hass, BACKUP_DOMAIN, {}) await setup_integration(hass, mock_config_entry) await hass.async_block_till_done() yield async def test_suggested_filenames() -> None: """Test the suggested_filenames function.""" backup = AgentBackup( backup_id="a1b2c3", date="2021-01-01T01:02:03+00:00", addons=[], database_included=False, extra_metadata={}, folders=[], homeassistant_included=False, homeassistant_version=None, name="my_pretty_backup", protected=False, size=0, ) tar_filename, metadata_filename = suggested_filenames(backup) assert tar_filename == "my_pretty_backup_2021-01-01_01.02_03000000.tar" assert ( metadata_filename == "my_pretty_backup_2021-01-01_01.02_03000000.metadata.json" ) async def test_agents_info( hass: HomeAssistant, hass_ws_client: WebSocketGenerator, mock_config_entry: MockConfigEntry, ) -> None: """Test backup agent info.""" client = await hass_ws_client(hass) await client.send_json_auto_id({"type": "backup/agents/info"}) response = await client.receive_json() assert response["success"] assert response["result"] == { "agents": [ {"agent_id": "backup.local", "name": "local"}, { "agent_id": f"{DOMAIN}.{mock_config_entry.entry_id}", "name": mock_config_entry.title, }, ], } async def test_agents_list_backups( hass: HomeAssistant, hass_ws_client: WebSocketGenerator, mock_config_entry: MockConfigEntry, mock_agent_backup: AgentBackup, ) -> None: """Test agent list backups.""" 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"] == [ { "addons": mock_agent_backup.addons, "agents": { f"{DOMAIN}.{mock_config_entry.entry_id}": { "protected": mock_agent_backup.protected, "size": mock_agent_backup.size, } }, "backup_id": mock_agent_backup.backup_id, "database_included": mock_agent_backup.database_included, "date": mock_agent_backup.date, "extra_metadata": mock_agent_backup.extra_metadata, "failed_addons": [], "failed_agent_ids": [], "failed_folders": [], "folders": mock_agent_backup.folders, "homeassistant_included": mock_agent_backup.homeassistant_included, "homeassistant_version": mock_agent_backup.homeassistant_version, "name": mock_agent_backup.name, "with_automatic_settings": None, } ] async def test_agents_get_backup( hass: HomeAssistant, hass_ws_client: WebSocketGenerator, mock_config_entry: MockConfigEntry, mock_agent_backup: AgentBackup, ) -> None: """Test agent get backup.""" client = await hass_ws_client(hass) await client.send_json_auto_id( {"type": "backup/details", "backup_id": mock_agent_backup.backup_id} ) response = await client.receive_json() assert response["success"] assert response["result"]["agent_errors"] == {} assert response["result"]["backup"] == { "addons": mock_agent_backup.addons, "agents": { f"{DOMAIN}.{mock_config_entry.entry_id}": { "protected": mock_agent_backup.protected, "size": mock_agent_backup.size, } }, "backup_id": mock_agent_backup.backup_id, "database_included": mock_agent_backup.database_included, "date": mock_agent_backup.date, "extra_metadata": mock_agent_backup.extra_metadata, "failed_addons": [], "failed_agent_ids": [], "failed_folders": [], "folders": mock_agent_backup.folders, "homeassistant_included": mock_agent_backup.homeassistant_included, "homeassistant_version": mock_agent_backup.homeassistant_version, "name": mock_agent_backup.name, "with_automatic_settings": None, } async def test_agents_get_backup_does_not_throw_on_not_found( hass: HomeAssistant, hass_ws_client: WebSocketGenerator, mock_client: MagicMock, ) -> None: """Test agent get backup does not throw on a backup not found.""" paginate = mock_client.get_paginator.return_value.paginate paginate.return_value.__aiter__.return_value = [{"Contents": []}] client = await hass_ws_client(hass) await client.send_json_auto_id({"type": "backup/details", "backup_id": "random"}) response = await client.receive_json() assert response["success"] assert response["result"]["agent_errors"] == {} assert response["result"]["backup"] is None async def test_agents_list_backups_with_corrupted_metadata( hass: HomeAssistant, mock_client: MagicMock, mock_config_entry: MockConfigEntry, caplog: pytest.LogCaptureFixture, mock_agent_backup: AgentBackup, ) -> None: """Test listing backups when one metadata file is corrupted.""" # Create agent agent = S3BackupAgent(hass, mock_config_entry) # Set up mock responses for both valid and corrupted metadata files paginate = mock_client.get_paginator.return_value.paginate paginate.return_value.__aiter__.return_value = [ { "Contents": [ { "Key": "valid_backup.metadata.json", "LastModified": "2023-01-01T00:00:00+00:00", }, { "Key": "corrupted_backup.metadata.json", "LastModified": "2023-01-01T00:00:00+00:00", }, ] } ] # Mock responses for get_object calls valid_metadata = json.dumps(mock_agent_backup.as_dict()) corrupted_metadata = "{invalid json content" async def mock_get_object(**kwargs): """Mock get_object with different responses based on the key.""" key = kwargs.get("Key", "") if "valid_backup" in key: mock_body = AsyncMock() mock_body.read.return_value = valid_metadata.encode() return {"Body": mock_body} # Corrupted metadata mock_body = AsyncMock() mock_body.read.return_value = corrupted_metadata.encode() return {"Body": mock_body} mock_client.get_object.side_effect = mock_get_object backups = await agent.async_list_backups() assert len(backups) == 1 assert backups[0].backup_id == mock_agent_backup.backup_id assert "Failed to process metadata file" in caplog.text async def test_agents_delete( hass: HomeAssistant, hass_ws_client: WebSocketGenerator, mock_client: MagicMock, ) -> None: """Test agent delete backup.""" client = await hass_ws_client(hass) await client.send_json_auto_id( { "type": "backup/delete", "backup_id": "23e64aec", } ) response = await client.receive_json() assert response["success"] assert response["result"] == {"agent_errors": {}} # Should delete both the tar and the metadata file assert mock_client.delete_object.call_count == 2 async def test_agents_delete_not_throwing_on_not_found( hass: HomeAssistant, hass_ws_client: WebSocketGenerator, mock_client: MagicMock, ) -> None: """Test agent delete backup does not throw on a backup not found.""" paginate = mock_client.get_paginator.return_value.paginate paginate.return_value.__aiter__.return_value = [{"Contents": []}] client = await hass_ws_client(hass) await client.send_json_auto_id( { "type": "backup/delete", "backup_id": "random", } ) response = await client.receive_json() assert response["success"] assert response["result"] == {"agent_errors": {}} assert mock_client.delete_object.call_count == 0 @pytest.mark.parametrize( "backup_size", [ 2**20, MULTIPART_MIN_PART_SIZE_BYTES, ], ids=["small", "large"], ) async def test_agents_upload_network_failure( hass_client: ClientSessionGenerator, caplog: pytest.LogCaptureFixture, mock_client: MagicMock, mock_config_entry: MockConfigEntry, mock_agent_backup: AgentBackup, ) -> None: """Test agent upload backup with network failure.""" client = await hass_client() with ( patch( "homeassistant.components.backup.manager.BackupManager.async_get_backup", return_value=mock_agent_backup, ), patch( "homeassistant.components.backup.manager.read_backup", return_value=mock_agent_backup, ), patch("pathlib.Path.open") as mocked_open, ): mocked_open.return_value.read = Mock(side_effect=[b"test", b""]) # simulate network failure mock_client.put_object.side_effect = mock_client.upload_part.side_effect = ( mock_client.abort_multipart_upload.side_effect ) = ConnectTimeoutError(endpoint_url=USER_INPUT[CONF_ENDPOINT_URL]) resp = await client.post( f"/api/backup/upload?agent_id={DOMAIN}.{mock_config_entry.entry_id}", data={"file": StringIO("test")}, ) assert resp.status == 201 assert "Upload failed for aws_s3" in caplog.text @pytest.mark.parametrize( "backup_size", [MULTIPART_MIN_PART_SIZE_BYTES * 2], ids=["large"] ) async def test_agents_upload_on_progress( hass: HomeAssistant, hass_client: ClientSessionGenerator, mock_client: MagicMock, mock_config_entry: MockConfigEntry, mock_agent_backup: AgentBackup, ) -> None: """Test agent upload backup emits UploadBackupEvent via on_progress.""" client = await hass_client() manager = hass.data[DATA_MANAGER] events: list[UploadBackupEvent] = [] def _collect(event: UploadBackupEvent) -> None: if isinstance(event, UploadBackupEvent): events.append(event) unsub = manager.async_subscribe_events(_collect) with ( patch( "homeassistant.components.backup.manager.BackupManager.async_get_backup", return_value=mock_agent_backup, ), patch( "homeassistant.components.backup.manager.read_backup", return_value=mock_agent_backup, ), patch("pathlib.Path.open") as mocked_open, ): mocked_open.return_value.read = Mock( side_effect=[ b"a" * mock_agent_backup.size, b"", ] ) resp = await client.post( f"/api/backup/upload?agent_id={DOMAIN}.{mock_config_entry.entry_id}", data={"file": StringIO("test")}, ) unsub() assert resp.status == 201 agent_id = f"{DOMAIN}.{mock_config_entry.entry_id}" agent_events = [e for e in events if e.agent_id == agent_id] assert len(agent_events) >= 2 assert all(e.total_bytes == mock_agent_backup.size for e in agent_events) # Verify events report distinct increasing byte counts uploaded_bytes = [e.uploaded_bytes for e in agent_events] assert uploaded_bytes == sorted(uploaded_bytes) assert len(set(uploaded_bytes)) == len(uploaded_bytes) # Verify at least one intermediate event (uploaded_bytes < total_bytes) assert agent_events[0].uploaded_bytes < agent_events[0].total_bytes async def test_agents_download( hass_client: ClientSessionGenerator, mock_client: MagicMock, mock_config_entry: MockConfigEntry, ) -> None: """Test agent download backup.""" client = await hass_client() backup_id = "23e64aec" resp = await client.get( f"/api/backup/download/{backup_id}?agent_id={DOMAIN}.{mock_config_entry.entry_id}" ) assert resp.status == 200 assert await resp.content.read() == b"backup data" # Coordinator first refresh reads metadata (1) + # download reads metadata (1) + tar (1) assert mock_client.get_object.call_count == 3 async def test_error_during_delete( hass: HomeAssistant, hass_ws_client: WebSocketGenerator, mock_client: MagicMock, mock_config_entry: MockConfigEntry, mock_agent_backup: AgentBackup, ) -> None: """Test the error wrapper.""" mock_client.delete_object.side_effect = BotoCoreError client = await hass_ws_client(hass) await client.send_json_auto_id( { "type": "backup/delete", "backup_id": mock_agent_backup.backup_id, } ) response = await client.receive_json() assert response["success"] assert response["result"] == { "agent_errors": { f"{DOMAIN}.{mock_config_entry.entry_id}": ( "Failed during async_delete_backup" ) } } async def test_cache_expiration( hass: HomeAssistant, mock_client: MagicMock, mock_agent_backup: AgentBackup, ) -> None: """Test that the cache expires correctly.""" # Mock the entry mock_entry = MockConfigEntry( domain=DOMAIN, data={"bucket": "test-bucket"}, unique_id="test-unique-id", title="Test S3", ) mock_entry.runtime_data = MagicMock(client=mock_client) # Create agent agent = S3BackupAgent(hass, mock_entry) # Reset call counts from coordinator's initial refresh mock_client.reset_mock() # Mock metadata response metadata_content = json.dumps(mock_agent_backup.as_dict()) mock_body = AsyncMock() mock_body.read.return_value = metadata_content.encode() paginate = mock_client.get_paginator.return_value.paginate paginate.return_value.__aiter__.return_value = [ { "Contents": [ { "Key": "test.metadata.json", "LastModified": "2023-01-01T00:00:00+00:00", } ] } ] mock_client.get_object.return_value = {"Body": mock_body} # First call should query S3 await agent.async_list_backups() assert mock_client.get_paginator.call_count == 1 assert mock_client.get_object.call_count == 1 # Second call should use cache await agent.async_list_backups() assert mock_client.get_paginator.call_count == 1 assert mock_client.get_object.call_count == 1 # Set cache to expire agent._cache_expiration = time() - 1 # Third call should query S3 again await agent.async_list_backups() assert mock_client.get_paginator.call_count == 2 assert mock_client.get_object.call_count == 2 async def test_listeners_get_cleaned_up(hass: HomeAssistant) -> None: """Test listener gets cleaned up.""" listener = MagicMock() remove_listener = async_register_backup_agents_listener(hass, listener=listener) hass.data[DATA_BACKUP_AGENT_LISTENERS] = [ listener ] # make sure it's the last listener remove_listener() assert DATA_BACKUP_AGENT_LISTENERS not in hass.data async def test_list_backups_with_pagination( hass: HomeAssistant, mock_config_entry: MockConfigEntry, ) -> None: """Test listing backups when paginating through multiple pages.""" # Create agent agent = S3BackupAgent(hass, mock_config_entry) # Create two different backups backup1 = AgentBackup( backup_id="backup1", date="2023-01-01T00:00:00+00:00", addons=[], database_included=False, extra_metadata={}, folders=[], homeassistant_included=False, homeassistant_version=None, name="Backup 1", protected=False, size=0, ) backup2 = AgentBackup( backup_id="backup2", date="2023-01-02T00:00:00+00:00", addons=[], database_included=False, extra_metadata={}, folders=[], homeassistant_included=False, homeassistant_version=None, name="Backup 2", protected=False, size=0, ) # Setup two pages of results page1 = { "Contents": [ { "Key": "backup1.metadata.json", "LastModified": "2023-01-01T00:00:00+00:00", }, {"Key": "backup1.tar", "LastModified": "2023-01-01T00:00:00+00:00"}, ] } page2 = { "Contents": [ { "Key": "backup2.metadata.json", "LastModified": "2023-01-02T00:00:00+00:00", }, {"Key": "backup2.tar", "LastModified": "2023-01-02T00:00:00+00:00"}, ] } # Setup mock client mock_client = mock_config_entry.runtime_data.client paginate = mock_client.get_paginator.return_value.paginate paginate.return_value.__aiter__.return_value = [ page1, page2, ] # Mock get_object responses based on the key async def mock_get_object(**kwargs): """Mock get_object with different responses based on the key.""" key = kwargs.get("Key", "") if "backup1" in key: mock_body = AsyncMock() mock_body.read.return_value = json.dumps(backup1.as_dict()).encode() return {"Body": mock_body} # backup2 mock_body = AsyncMock() mock_body.read.return_value = json.dumps(backup2.as_dict()).encode() return {"Body": mock_body} mock_client.get_object.side_effect = mock_get_object # List backups and verify we got both backups = await agent.async_list_backups() assert len(backups) == 2 backup_ids = {backup.backup_id for backup in backups} assert backup_ids == {"backup1", "backup2"} @pytest.mark.parametrize( ("config_entry_extra_data", "expected_paginate_extra_kwargs"), [ ({"prefix": "backups/home"}, {"Prefix": "backups/home/"}), ({}, {}), ], ids=["with_prefix", "no_prefix"], ) async def test_agent_list_backups_parametrized( hass: HomeAssistant, hass_ws_client: WebSocketGenerator, mock_config_entry: MockConfigEntry, mock_client: MagicMock, mock_agent_backup: AgentBackup, config_entry_extra_data: dict, expected_paginate_extra_kwargs: dict, ) -> None: """Test agent list backups with and without prefix.""" 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"] == {} # Verify pagination call with expected parameters mock_client.get_paginator.return_value.paginate.assert_called_with( **{"Bucket": "test"} | expected_paginate_extra_kwargs ) @pytest.mark.parametrize( ("config_entry_extra_data", "expected_key_prefix"), [ ({"prefix": "backups/home"}, "backups/home/"), ({}, ""), ], ids=["with_prefix", "no_prefix"], ) async def test_agent_delete_backup_parametrized( hass: HomeAssistant, hass_ws_client: WebSocketGenerator, mock_client: MagicMock, mock_config_entry: MockConfigEntry, mock_agent_backup: AgentBackup, expected_key_prefix: str, ) -> None: """Test agent delete backup with and without prefix.""" client = await hass_ws_client(hass) await client.send_json_auto_id( { "type": "backup/delete", "backup_id": "23e64aec", } ) response = await client.receive_json() assert response["success"] assert response["result"] == {"agent_errors": {}} tar_filename, metadata_filename = suggested_filenames(mock_agent_backup) expected_tar_key = f"{expected_key_prefix}{tar_filename}" expected_metadata_key = f"{expected_key_prefix}{metadata_filename}" mock_client.delete_object.assert_any_call(Bucket="test", Key=expected_tar_key) mock_client.delete_object.assert_any_call(Bucket="test", Key=expected_metadata_key) async def _upload_backup( hass_client: ClientSessionGenerator, agent_id: str, mock_agent_backup: AgentBackup, ) -> None: """Perform a backup upload with the necessary mocks set up.""" client = await hass_client() with ( patch( "homeassistant.components.backup.manager.BackupManager.async_get_backup", return_value=mock_agent_backup, ), patch( "homeassistant.components.backup.manager.read_backup", return_value=mock_agent_backup, ), patch("pathlib.Path.open") as mocked_open, ): # we must emit at least two chunks # the "appendix" chunk triggers the upload of the final buffer part mocked_open.return_value.read = Mock( side_effect=[ b"a" * mock_agent_backup.size, b"appendix", b"", ] ) resp = await client.post( f"/api/backup/upload?agent_id={agent_id}", data={"file": StringIO("test")}, ) assert resp.status == 201 @pytest.mark.parametrize( ("config_entry_extra_data", "expected_key_prefix"), [ ({"prefix": "backups/home"}, "backups/home/"), ({}, ""), ], ids=["with_prefix", "no_prefix"], ) async def test_agent_upload_small_backup_parametrized( hass_client: ClientSessionGenerator, caplog: pytest.LogCaptureFixture, mock_client: MagicMock, mock_config_entry: MockConfigEntry, mock_agent_backup: AgentBackup, expected_key_prefix: str, ) -> None: """Test agent upload small backup with and without prefix.""" await _upload_backup( hass_client, f"{DOMAIN}.{mock_config_entry.entry_id}", mock_agent_backup ) assert f"Uploading backup {mock_agent_backup.backup_id}" in caplog.text assert mock_client.create_multipart_upload.await_count == 0 assert mock_client.upload_part.await_count == 0 assert mock_client.complete_multipart_upload.await_count == 0 assert mock_client.put_object.await_count == 2 tar_filename, metadata_filename = suggested_filenames(mock_agent_backup) mock_client.put_object.assert_has_calls( [ call(Bucket="test", Key=f"{expected_key_prefix}{tar_filename}", Body=ANY), call( Bucket="test", Key=f"{expected_key_prefix}{metadata_filename}", Body=ANY, ), ] ) @pytest.mark.parametrize("backup_size", [MULTIPART_MIN_PART_SIZE_BYTES], ids=["large"]) @pytest.mark.parametrize( ("config_entry_extra_data", "expected_key_prefix"), [ ({"prefix": "backups/home"}, "backups/home/"), ({}, ""), ], ids=["with_prefix", "no_prefix"], ) async def test_agent_upload_large_backup_parametrized( hass_client: ClientSessionGenerator, caplog: pytest.LogCaptureFixture, mock_client: MagicMock, mock_config_entry: MockConfigEntry, mock_agent_backup: AgentBackup, expected_key_prefix: str, ) -> None: """Test agent upload large (multipart) backup with and without prefix.""" await _upload_backup( hass_client, f"{DOMAIN}.{mock_config_entry.entry_id}", mock_agent_backup ) tar_filename, metadata_filename = suggested_filenames(mock_agent_backup) tar_key = f"{expected_key_prefix}{tar_filename}" metadata_key = f"{expected_key_prefix}{metadata_filename}" assert f"Uploading backup {mock_agent_backup.backup_id}" in caplog.text assert mock_client.create_multipart_upload.await_count == 1 assert mock_client.upload_part.await_count == 2 assert mock_client.complete_multipart_upload.await_count == 1 assert mock_client.put_object.await_count == 1 mock_client.create_multipart_upload.assert_called_with(Bucket="test", Key=tar_key) mock_client.upload_part.assert_has_calls( [ call( Bucket="test", Key=tar_key, PartNumber=1, UploadId="upload_id", Body=ANY, ), call( Bucket="test", Key=tar_key, PartNumber=2, UploadId="upload_id", Body=ANY, ), ] ) mock_client.complete_multipart_upload.assert_called_with( Bucket="test", Key=tar_key, UploadId="upload_id", MultipartUpload=ANY, ) mock_client.put_object.assert_called_with(Bucket="test", Key=metadata_key, Body=ANY) @pytest.mark.parametrize( ("config_entry_extra_data", "expected_key_prefix"), [ ({"prefix": "backups/home"}, "backups/home/"), ({}, ""), ], ids=["with_prefix", "no_prefix"], ) async def test_agent_download_backup_parametrized( hass: HomeAssistant, hass_client: ClientSessionGenerator, mock_client: MagicMock, mock_config_entry: MockConfigEntry, mock_agent_backup: AgentBackup, expected_key_prefix: str, ) -> None: """Test agent download backup with and without prefix.""" client = await hass_client() backup_id = "23e64aec" resp = await client.get( f"/api/backup/download/{backup_id}?agent_id={DOMAIN}.{mock_config_entry.entry_id}" ) assert resp.status == 200 assert await resp.content.read() == b"backup data" tar_filename, _ = suggested_filenames(mock_agent_backup) expected_tar_key = f"{expected_key_prefix}{tar_filename}" mock_client.get_object.assert_any_call(Bucket="test", Key=expected_tar_key)