Skip to content

Commit

Permalink
Use test helper for creating a mocked backup agent in backup tests (#…
Browse files Browse the repository at this point in the history 
…138312)

* Use test helper for creating a mocked backup agent in backup tests

* Adjust according to discussion
  • Loading branch information
@emontnemery
emontnemery authored Feb 12, 2025
1 parent 1393f41 commit a6c5144
Show file tree
Hide file tree
Showing 6 changed files with 191 additions and 239 deletions.
130 changes: 39 additions & 91 deletions tests/components/backup/common.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,6 @@
)
from homeassistant.components.backup.const import DATA_MANAGER
from homeassistant.core import HomeAssistant
from homeassistant.helpers.typing import ConfigType
from homeassistant.setup import async_setup_component

from tests.common import MockPlatform, mock_platform
Expand Down Expand Up @@ -64,95 +63,45 @@ async def aiter_from_iter(iterable: Iterable) -> AsyncIterator:
yield i


class BackupAgentTest(BackupAgent):
"""Test backup agent."""

domain = "test"

def __init__(self, name: str, backups: list[AgentBackup] | None = None) -> None:
"""Initialize the backup agent."""
self.name = name
self.unique_id = name
if backups is None:
backups = [
AgentBackup(
addons=[AddonInfo(name="Test", slug="test", version="1.0.0")],
backup_id="abc123",
database_included=True,
date="1970-01-01T00:00:00Z",
extra_metadata={},
folders=[Folder.MEDIA, Folder.SHARE],
homeassistant_included=True,
homeassistant_version="2024.12.0",
name="Test",
protected=False,
size=13,
)
]

self._backup_data: bytearray | None = None
self._backups = {backup.backup_id: backup for backup in backups}

async def async_download_backup(
self,
backup_id: str,
**kwargs: Any,
) -> AsyncIterator[bytes]:
"""Download a backup file."""
return AsyncMock(spec_set=["__aiter__"])
def mock_backup_agent(name: str, backups: list[AgentBackup] | None = None) -> Mock:
"""Create a mock backup agent."""

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

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

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

async def async_list_backups(self, **kwargs: Any) -> list[AgentBackup]:
"""List backups."""
return list(self._backups.values())

async def async_get_backup(
self,
backup_id: str,
**kwargs: Any,
) -> AgentBackup | None:
"""Return a backup."""
return self._backups.get(backup_id)

async def async_delete_backup(
self,
backup_id: str,
**kwargs: Any,
) -> None:
"""Delete a backup file."""


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

async def get_backup(backup_id: str, **kwargs: Any) -> AgentBackup | None:
"""Get a backup."""
return next((b for b in backups if b.backup_id == backup_id), None)
backup_data += chunk
backups_data[backup.backup_id] = backup_data

backups = backups or []
backups_data: dict[str, bytes] = {}
mock_agent = Mock(spec=BackupAgent)
mock_agent.domain = "test"
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(
spec_set=[BackupAgent.async_delete_backup]
)
mock_agent.async_download_backup = AsyncMock(
side_effect=BackupNotFound, spec_set=[BackupAgent.async_download_backup]
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]
Expand All @@ -161,57 +110,56 @@ async def get_backup(backup_id: str, **kwargs: Any) -> AgentBackup | None:
return_value=backups, spec_set=[BackupAgent.async_list_backups]
)
mock_agent.async_upload_backup = AsyncMock(
spec_set=[BackupAgent.async_upload_backup]
side_effect=upload_backup,
spec_set=[BackupAgent.async_upload_backup],
)
return mock_agent


async def setup_backup_integration(
hass: HomeAssistant,
with_hassio: bool = False,
configuration: ConfigType | None = None,
*,
backups: dict[str, list[AgentBackup]] | None = None,
remote_agents: list[str] | None = None,
) -> bool:
) -> 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
),
):
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))
platform = Mock(
async_get_backup_agents=AsyncMock(
return_value=[BackupAgentTest(agent, []) for agent in remote_agents]
return_value=list(remote_agents_dict.values())
),
spec_set=BackupAgentPlatformProtocol,
)

mock_platform(hass, f"{TEST_DOMAIN}.backup", platform or MockPlatform())
assert await async_setup_component(hass, TEST_DOMAIN, {})

result = await async_setup_component(hass, DOMAIN, configuration or {})
assert await async_setup_component(hass, DOMAIN, {})
await hass.async_block_till_done()
if not backups:
return result

for agent_id, agent_backups in backups.items():
if with_hassio and agent_id == LOCAL_AGENT_ID:
continue
agent = hass.data[DATA_MANAGER].backup_agents[agent_id]
if LOCAL_AGENT_ID not in backups or with_hassio:
return remote_agents_dict

async def open_stream() -> AsyncIterator[bytes]:
"""Open a stream."""
return aiter_from_iter((b"backup data",))
agent = hass.data[DATA_MANAGER].backup_agents[LOCAL_AGENT_ID]

for backup in agent_backups:
await agent.async_upload_backup(open_stream=open_stream, backup=backup)
if agent_id == LOCAL_AGENT_ID:
agent._loaded_backups = True
for backup in backups[LOCAL_AGENT_ID]:
await agent.async_upload_backup(open_stream=None, backup=backup)
agent._loaded_backups = True

return result
return remote_agents_dict


async def setup_backup_platform(
Expand Down
48 changes: 30 additions & 18 deletions tests/components/backup/snapshots/test_websocket.ambr
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3482,13 +3482,15 @@
'agents': dict({
'domain.test': dict({
'protected': False,
'size': 13,
'size': 0,
}),
}),
'backup_id': 'abc123',
'database_included': True,
'date': '1970-01-01T00:00:00Z',
'date': '1970-01-01T00:00:00.000Z',
'extra_metadata': dict({
'instance_id': 'our_uuid',
'with_automatic_settings': True,
}),
'failed_agent_ids': list([
]),
Expand All @@ -3499,7 +3501,7 @@
'homeassistant_included': True,
'homeassistant_version': '2024.12.0',
'name': 'Test',
'with_automatic_settings': None,
'with_automatic_settings': True,
}),
]),
'last_attempted_automatic_backup': None,
Expand Down Expand Up @@ -3543,13 +3545,15 @@
'agents': dict({
'domain.test': dict({
'protected': False,
'size': 13,
'size': 0,
}),
}),
'backup_id': 'abc123',
'database_included': True,
'date': '1970-01-01T00:00:00Z',
'date': '1970-01-01T00:00:00.000Z',
'extra_metadata': dict({
'instance_id': 'our_uuid',
'with_automatic_settings': True,
}),
'failed_agent_ids': list([
'test.remote',
Expand All @@ -3561,7 +3565,7 @@
'homeassistant_included': True,
'homeassistant_version': '2024.12.0',
'name': 'Test',
'with_automatic_settings': None,
'with_automatic_settings': True,
}),
]),
'last_attempted_automatic_backup': None,
Expand Down Expand Up @@ -3604,13 +3608,15 @@
'agents': dict({
'domain.test': dict({
'protected': False,
'size': 13,
'size': 0,
}),
}),
'backup_id': 'abc123',
'database_included': True,
'date': '1970-01-01T00:00:00Z',
'date': '1970-01-01T00:00:00.000Z',
'extra_metadata': dict({
'instance_id': 'our_uuid',
'with_automatic_settings': True,
}),
'failed_agent_ids': list([
]),
Expand All @@ -3621,7 +3627,7 @@
'homeassistant_included': True,
'homeassistant_version': '2024.12.0',
'name': 'Test',
'with_automatic_settings': None,
'with_automatic_settings': True,
}),
]),
'last_attempted_automatic_backup': None,
Expand Down Expand Up @@ -3664,13 +3670,15 @@
'agents': dict({
'domain.test': dict({
'protected': False,
'size': 13,
'size': 0,
}),
}),
'backup_id': 'abc123',
'database_included': True,
'date': '1970-01-01T00:00:00Z',
'date': '1970-01-01T00:00:00.000Z',
'extra_metadata': dict({
'instance_id': 'our_uuid',
'with_automatic_settings': True,
}),
'failed_agent_ids': list([
]),
Expand All @@ -3681,7 +3689,7 @@
'homeassistant_included': True,
'homeassistant_version': '2024.12.0',
'name': 'Test',
'with_automatic_settings': None,
'with_automatic_settings': True,
}),
]),
'last_attempted_automatic_backup': None,
Expand Down Expand Up @@ -3725,13 +3733,15 @@
'agents': dict({
'domain.test': dict({
'protected': False,
'size': 13,
'size': 0,
}),
}),
'backup_id': 'abc123',
'database_included': True,
'date': '1970-01-01T00:00:00Z',
'date': '1970-01-01T00:00:00.000Z',
'extra_metadata': dict({
'instance_id': 'our_uuid',
'with_automatic_settings': True,
}),
'failed_agent_ids': list([
]),
Expand All @@ -3742,7 +3752,7 @@
'homeassistant_included': True,
'homeassistant_version': '2024.12.0',
'name': 'Test',
'with_automatic_settings': None,
'with_automatic_settings': True,
}),
]),
'last_attempted_automatic_backup': None,
Expand Down Expand Up @@ -3786,13 +3796,15 @@
'agents': dict({
'domain.test': dict({
'protected': False,
'size': 13,
'size': 0,
}),
}),
'backup_id': 'abc123',
'database_included': True,
'date': '1970-01-01T00:00:00Z',
'date': '1970-01-01T00:00:00.000Z',
'extra_metadata': dict({
'instance_id': 'our_uuid',
'with_automatic_settings': True,
}),
'failed_agent_ids': list([
'test.remote',
Expand All @@ -3804,7 +3816,7 @@
'homeassistant_included': True,
'homeassistant_version': '2024.12.0',
'name': 'Test',
'with_automatic_settings': None,
'with_automatic_settings': True,
}),
]),
'last_attempted_automatic_backup': None,
Expand Down
14 changes: 5 additions & 9 deletions tests/components/backup/test_http.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,6 @@

from .common import (
TEST_BACKUP_ABC123,
BackupAgentTest,
aiter_from_iter,
mock_backup_agent,
setup_backup_integration,
Expand Down Expand Up @@ -65,19 +64,16 @@ async def test_downloading_remote_backup(
hass_client: ClientSessionGenerator,
) -> None:
"""Test downloading a remote backup."""

await setup_backup_integration(
hass, backups={"test.test": [TEST_BACKUP_ABC123]}, remote_agents=["test"]
hass, backups={"test.test": [TEST_BACKUP_ABC123]}, remote_agents=["test.test"]
)

client = await hass_client()

with (
patch.object(BackupAgentTest, "async_download_backup") as download_mock,
):
download_mock.return_value.__aiter__.return_value = iter((b"backup data",))
resp = await client.get("/api/backup/download/abc123?agent_id=test.test")
assert resp.status == 200
assert await resp.content.read() == b"backup data"
resp = await client.get("/api/backup/download/abc123?agent_id=test.test")
assert resp.status == 200
assert await resp.content.read() == b"backup data"


async def test_downloading_local_encrypted_backup_file_not_found(
Expand Down
Loading

0 comments on commit a6c5144

Please sign in to comment.