core/tests/components/backup/common.py

"""Common helpers for the Backup integration tests."""

from collections.abc import AsyncIterator, Buffer, Callable, Coroutine, Iterable
from pathlib import Path
from typing import Any, cast
from unittest.mock import AsyncMock, Mock, patch

from homeassistant.components.backup import (
    DOMAIN,
    AddonInfo,
    AgentBackup,
    BackupAgent,
    BackupAgentPlatformProtocol,
    BackupNotFound,
    Folder,
)
from homeassistant.components.backup.backup import CoreLocalBackupAgent
from homeassistant.components.backup.const import DATA_MANAGER
from homeassistant.core import HomeAssistant
from homeassistant.setup import async_setup_component

from tests.common import mock_platform

LOCAL_AGENT_ID = f"{DOMAIN}.local"

TEST_BACKUP_ABC123 = AgentBackup(
    addons=[AddonInfo(name="Test", slug="test", version="1.0.0")],
    backup_id="abc123",
    database_included=True,
    date="1970-01-01T00:00:00.000Z",
    extra_metadata={"instance_id": "our_uuid", "with_automatic_settings": True},
    folders=[Folder.MEDIA, Folder.SHARE],
    homeassistant_included=True,
    homeassistant_version="2024.12.0",
    name="Test",
    protected=False,
    size=0,
)
TEST_BACKUP_PATH_ABC123 = Path("abc123.tar")

TEST_BACKUP_DEF456 = AgentBackup(
    addons=[],
    backup_id="def456",
    database_included=False,
    date="1980-01-01T00:00:00.000Z",
    extra_metadata={"instance_id": "unknown_uuid", "with_automatic_settings": True},
    folders=[Folder.MEDIA, Folder.SHARE],
    homeassistant_included=True,
    homeassistant_version="2024.12.0",
    name="Test 2",
    protected=False,
    size=1,
)
TEST_BACKUP_PATH_DEF456 = Path("custom_def456.tar")

TEST_DOMAIN = "test"


async def aiter_from_iter(iterable: Iterable) -> AsyncIterator:
    """Convert an iterable to an async iterator."""
    for i in iterable:
        yield i


def mock_backup_agent(name: str, backups: list[AgentBackup] | None = None) -> Mock:
    """Create a mock backup agent."""

    async def delete_backup(backup_id: str, **kwargs: Any) -> None:
        """Mock delete."""
        await get_backup(backup_id)

    async def download_backup(backup_id: str, **kwargs: Any) -> AsyncIterator[bytes]:
        """Mock download."""
        return aiter_from_iter((backups_data.get(backup_id, b"backup data"),))

    async def get_backup(backup_id: str, **kwargs: Any) -> AgentBackup:
        """Get a backup."""
        backup = next((b for b in _backups if b.backup_id == backup_id), None)
        if backup is None:
            raise BackupNotFound
        return backup

    async def upload_backup(
        *,
        open_stream: Callable[[], Coroutine[Any, Any, AsyncIterator[bytes]]],
        backup: AgentBackup,
        **kwargs: Any,
    ) -> None:
        """Upload a backup."""
        _backups.append(backup)
        backup_stream = await open_stream()
        backup_data = bytearray()
        async for chunk in backup_stream:
            backup_data += chunk
        backups_data[backup.backup_id] = backup_data

    _backups = backups or []
    backups_data: dict[str, Buffer] = {}
    mock_agent = Mock(spec=BackupAgent)
    mock_agent.domain = TEST_DOMAIN
    mock_agent.name = name
    mock_agent.unique_id = name
    type(mock_agent).agent_id = BackupAgent.agent_id
    mock_agent.async_delete_backup = AsyncMock(
        side_effect=delete_backup, spec_set=[BackupAgent.async_delete_backup]
    )
    mock_agent.async_download_backup = AsyncMock(
        side_effect=download_backup, spec_set=[BackupAgent.async_download_backup]
    )
    mock_agent.async_get_backup = AsyncMock(
        side_effect=get_backup, spec_set=[BackupAgent.async_get_backup]
    )
    mock_agent.async_list_backups = AsyncMock(
        return_value=_backups, spec_set=[BackupAgent.async_list_backups]
    )
    mock_agent.async_upload_backup = AsyncMock(
        side_effect=upload_backup,
        spec_set=[BackupAgent.async_upload_backup],
    )
    return mock_agent


async def setup_backup_integration(
    hass: HomeAssistant,
    with_hassio: bool = False,
    *,
    backups: dict[str, list[AgentBackup]] | None = None,
    remote_agents: list[str] | None = None,
) -> dict[str, Mock]:
    """Set up the Backup integration."""
    backups = backups or {}
    with (
        patch("homeassistant.components.backup.is_hassio", return_value=with_hassio),
        patch(
            "homeassistant.components.backup.backup.is_hassio", return_value=with_hassio
        ),
        patch(
            "homeassistant.components.backup.services.is_hassio",
            return_value=with_hassio,
        ),
    ):
        remote_agents = remote_agents or []
        remote_agents_dict = {}
        for agent in remote_agents:
            if not agent.startswith(f"{TEST_DOMAIN}."):
                raise ValueError(f"Invalid agent_id: {agent}")
            name = agent.partition(".")[2]
            remote_agents_dict[agent] = mock_backup_agent(name, backups.get(agent))
        if remote_agents:
            platform = Mock(
                async_get_backup_agents=AsyncMock(
                    return_value=list(remote_agents_dict.values())
                ),
                spec_set=BackupAgentPlatformProtocol,
            )
            await setup_backup_platform(hass, domain=TEST_DOMAIN, platform=platform)

        assert await async_setup_component(hass, DOMAIN, {})
        await hass.async_block_till_done()

        if LOCAL_AGENT_ID not in backups or with_hassio:
            return remote_agents_dict

        local_agent = cast(
            CoreLocalBackupAgent, hass.data[DATA_MANAGER].backup_agents[LOCAL_AGENT_ID]
        )

        for backup in backups[LOCAL_AGENT_ID]:
            await local_agent.async_upload_backup(
                open_stream=AsyncMock(
                    side_effect=RuntimeError("Local agent does not open stream")
                ),
                backup=backup,
                on_progress=lambda *, on_progress, **_: None,
            )
        local_agent._loaded_backups = True

        return remote_agents_dict


async def setup_backup_platform(
    hass: HomeAssistant,
    *,
    domain: str,
    platform: Any,
) -> None:
    """Set up a mock domain."""
    mock_platform(hass, f"{domain}.backup", platform)
    assert await async_setup_component(hass, domain, {})
    await hass.async_block_till_done()