From 9c93945b0e0e0442436e3535e7f2a2f7301a8851 Mon Sep 17 00:00:00 2001 From: tronikos Date: Fri, 3 Jan 2025 20:21:13 +0000 Subject: [PATCH 01/23] Add Google Drive integration for backup --- CODEOWNERS | 2 + homeassistant/brands/google.json | 1 + .../components/google_drive/__init__.py | 45 +++ homeassistant/components/google_drive/api.py | 71 +++++ .../google_drive/application_credentials.py | 21 ++ .../components/google_drive/backup.py | 272 ++++++++++++++++++ .../components/google_drive/config_flow.py | 105 +++++++ .../components/google_drive/const.py | 11 + .../components/google_drive/manifest.json | 14 + .../google_drive/quality_scale.yaml | 113 ++++++++ .../components/google_drive/strings.json | 36 +++ .../generated/application_credentials.py | 1 + homeassistant/generated/config_flows.py | 1 + homeassistant/generated/integrations.json | 6 + requirements_all.txt | 3 + requirements_test_all.txt | 3 + tests/components/google_drive/__init__.py | 1 + 17 files changed, 706 insertions(+) create mode 100644 homeassistant/components/google_drive/__init__.py create mode 100644 homeassistant/components/google_drive/api.py create mode 100644 homeassistant/components/google_drive/application_credentials.py create mode 100644 homeassistant/components/google_drive/backup.py create mode 100644 homeassistant/components/google_drive/config_flow.py create mode 100644 homeassistant/components/google_drive/const.py create mode 100644 homeassistant/components/google_drive/manifest.json create mode 100644 homeassistant/components/google_drive/quality_scale.yaml create mode 100644 homeassistant/components/google_drive/strings.json create mode 100644 tests/components/google_drive/__init__.py diff --git a/CODEOWNERS b/CODEOWNERS index f43cdf457c87be..388110491cf1e4 100644 --- a/CODEOWNERS +++ b/CODEOWNERS @@ -566,6 +566,8 @@ build.json @home-assistant/supervisor /tests/components/google_assistant_sdk/ @tronikos /homeassistant/components/google_cloud/ @lufton @tronikos /tests/components/google_cloud/ @lufton @tronikos +/homeassistant/components/google_drive/ @tronikos +/tests/components/google_drive/ @tronikos /homeassistant/components/google_generative_ai_conversation/ @tronikos /tests/components/google_generative_ai_conversation/ @tronikos /homeassistant/components/google_mail/ @tkdrob diff --git a/homeassistant/brands/google.json b/homeassistant/brands/google.json index 028fa544a5f742..872cfc0aac5021 100644 --- a/homeassistant/brands/google.json +++ b/homeassistant/brands/google.json @@ -5,6 +5,7 @@ "google_assistant", "google_assistant_sdk", "google_cloud", + "google_drive", "google_generative_ai_conversation", "google_mail", "google_maps", diff --git a/homeassistant/components/google_drive/__init__.py b/homeassistant/components/google_drive/__init__.py new file mode 100644 index 00000000000000..b3ca0e932e6611 --- /dev/null +++ b/homeassistant/components/google_drive/__init__.py @@ -0,0 +1,45 @@ +"""The Google Drive integration.""" + +from __future__ import annotations + +from collections.abc import Callable + +from homeassistant.config_entries import ConfigEntry +from homeassistant.core import HomeAssistant +from homeassistant.helpers.config_entry_oauth2_flow import ( + OAuth2Session, + async_get_config_entry_implementation, +) +from homeassistant.util.hass_dict import HassKey + +from .api import AsyncConfigEntryAuth +from .const import DOMAIN + +DATA_BACKUP_AGENT_LISTENERS: HassKey[list[Callable[[], None]]] = HassKey( + f"{DOMAIN}.backup_agent_listeners" +) + +type GoogleDriveConfigEntry = ConfigEntry[AsyncConfigEntryAuth] + + +async def async_setup_entry(hass: HomeAssistant, entry: GoogleDriveConfigEntry) -> bool: + """Set up Google Drive from a config entry.""" + implementation = await async_get_config_entry_implementation(hass, entry) + session = OAuth2Session(hass, entry, implementation) + auth = AsyncConfigEntryAuth(hass, session) + await auth.check_and_refresh_token() + entry.runtime_data = auth + return True + + +async def async_unload_entry( + hass: HomeAssistant, entry: GoogleDriveConfigEntry +) -> bool: + """Unload a config entry.""" + hass.async_create_task(_notify_backup_listeners(hass), eager_start=False) + return True + + +async def _notify_backup_listeners(hass: HomeAssistant) -> None: + for listener in hass.data.get(DATA_BACKUP_AGENT_LISTENERS, []): + listener() diff --git a/homeassistant/components/google_drive/api.py b/homeassistant/components/google_drive/api.py new file mode 100644 index 00000000000000..80f623fead67d6 --- /dev/null +++ b/homeassistant/components/google_drive/api.py @@ -0,0 +1,71 @@ +"""API for Google Drive bound to Home Assistant OAuth.""" + +from __future__ import annotations + +from datetime import datetime +from typing import Any + +from aiogoogle.auth import UserCreds +from aiohttp.client_exceptions import ClientError, ClientResponseError +from google.auth.exceptions import RefreshError + +from homeassistant.config_entries import ConfigEntryState +from homeassistant.const import CONF_ACCESS_TOKEN +from homeassistant.core import HomeAssistant +from homeassistant.exceptions import ( + ConfigEntryAuthFailed, + ConfigEntryNotReady, + HomeAssistantError, +) +from homeassistant.helpers import config_entry_oauth2_flow + + +def convert_to_user_creds(token: dict[str, Any]) -> UserCreds: + """Convert an OAuth2Session token to UserCreds.""" + return UserCreds( + access_token=token["access_token"], + expires_at=datetime.fromtimestamp(token["expires_at"]).isoformat(), + ) + + +class AsyncConfigEntryAuth: + """Provide Google Drive authentication tied to an OAuth2 based config entry.""" + + def __init__( + self, + hass: HomeAssistant, + oauth2_session: config_entry_oauth2_flow.OAuth2Session, + ) -> None: + """Initialize Google Drive Auth.""" + self._hass = hass + self.oauth_session = oauth2_session + + @property + def access_token(self) -> str: + """Return the access token.""" + return self.oauth_session.token[CONF_ACCESS_TOKEN] + + async def check_and_refresh_token(self) -> str: + """Check the token.""" + try: + await self.oauth_session.async_ensure_token_valid() + except (RefreshError, ClientResponseError, ClientError) as ex: + if ( + self.oauth_session.config_entry.state + is ConfigEntryState.SETUP_IN_PROGRESS + ): + if isinstance(ex, ClientResponseError) and 400 <= ex.status < 500: + raise ConfigEntryAuthFailed( + "OAuth session is not valid, reauth required" + ) from ex + raise ConfigEntryNotReady from ex + if ( + isinstance(ex, RefreshError) + or hasattr(ex, "status") + and ex.status == 400 + ): + self.oauth_session.config_entry.async_start_reauth( + self.oauth_session.hass + ) + raise HomeAssistantError(ex) from ex + return self.access_token diff --git a/homeassistant/components/google_drive/application_credentials.py b/homeassistant/components/google_drive/application_credentials.py new file mode 100644 index 00000000000000..c2f59b298cbeca --- /dev/null +++ b/homeassistant/components/google_drive/application_credentials.py @@ -0,0 +1,21 @@ +"""application_credentials platform for Google Drive.""" + +from homeassistant.components.application_credentials import AuthorizationServer +from homeassistant.core import HomeAssistant + + +async def async_get_authorization_server(hass: HomeAssistant) -> AuthorizationServer: + """Return authorization server.""" + return AuthorizationServer( + "https://accounts.google.com/o/oauth2/v2/auth", + "https://oauth2.googleapis.com/token", + ) + + +async def async_get_description_placeholders(hass: HomeAssistant) -> dict[str, str]: + """Return description placeholders for the credentials dialog.""" + return { + "oauth_consent_url": "https://console.cloud.google.com/apis/credentials/consent", + "more_info_url": "https://www.home-assistant.io/integrations/google_drive/", + "oauth_creds_url": "https://console.cloud.google.com/apis/credentials", + } diff --git a/homeassistant/components/google_drive/backup.py b/homeassistant/components/google_drive/backup.py new file mode 100644 index 00000000000000..44b7f1df9abe36 --- /dev/null +++ b/homeassistant/components/google_drive/backup.py @@ -0,0 +1,272 @@ +"""Backup platform for the Google Drive integration.""" + +from __future__ import annotations + +from collections.abc import AsyncIterator, Callable, Coroutine +import json +from typing import Any, Self + +from aiogoogle import Aiogoogle +from aiogoogle.auth import UserCreds +from aiogoogle.excs import AiogoogleError +from aiohttp import ClientError, StreamReader + +from homeassistant.components.backup import AgentBackup, BackupAgent, BackupAgentError +from homeassistant.core import HomeAssistant, callback +from homeassistant.exceptions import HomeAssistantError +from homeassistant.helpers.aiohttp_client import async_get_clientsession + +from . import DATA_BACKUP_AGENT_LISTENERS, GoogleDriveConfigEntry +from .api import AsyncConfigEntryAuth, convert_to_user_creds +from .const import DOMAIN + +# Google Drive only supports string key value pairs as properties. +# We convert any other fields to JSON strings. +_AGENT_BACKUP_SIMPLE_FIELDS = [ + "backup_id", + "date", + "database_included", + "homeassistant_included", + "homeassistant_version", + "name", + "protected", + "size", +] + + +def _convert_agent_backup_to_properties(backup: AgentBackup) -> dict[str, str]: + return { + k: v if k in _AGENT_BACKUP_SIMPLE_FIELDS else json.dumps(v) + for k, v in backup.as_dict().items() + } + + +def _convert_properties_to_agent_backup(d: dict[str, str]) -> AgentBackup: + return AgentBackup.from_dict( + { + k: v if k in _AGENT_BACKUP_SIMPLE_FIELDS else json.loads(v) + for k, v in d.items() + } + ) + + +async def async_get_backup_agents( + hass: HomeAssistant, + **kwargs: Any, +) -> list[BackupAgent]: + """Return a list of backup agents.""" + return [ + GoogleDriveBackupAgent(hass=hass, config_entry=config_entry) + for config_entry in hass.config_entries.async_loaded_entries(DOMAIN) + ] + + +@callback +def async_register_backup_agents_listener( + hass: HomeAssistant, + *, + listener: Callable[[], None], + **kwargs: Any, +) -> Callable[[], None]: + """Register a listener to be called when agents are added or removed. + + :return: A function to unregister the listener. + """ + hass.data.setdefault(DATA_BACKUP_AGENT_LISTENERS, []).append(listener) + + @callback + def remove_listener() -> None: + """Remove the listener.""" + hass.data[DATA_BACKUP_AGENT_LISTENERS].remove(listener) + + return remove_listener + + +class ChunkAsyncStreamIterator: + """Async iterator for chunked streams. + + Based on aiohttp.streams.ChunkTupleAsyncStreamIterator, but yields + bytes instead of tuple[bytes, bool]. + """ + + __slots__ = ("_stream",) + + def __init__(self, stream: StreamReader) -> None: + """Initialize.""" + self._stream = stream + + def __aiter__(self) -> Self: + """Iterate.""" + return self + + async def __anext__(self) -> bytes: + """Yield next chunk.""" + rv = await self._stream.readchunk() + if rv == (b"", False): + raise StopAsyncIteration + return rv[0] + + +class GoogleDriveBackupAgent(BackupAgent): + """Google Drive backup agent.""" + + domain = DOMAIN + + def __init__( + self, hass: HomeAssistant, config_entry: GoogleDriveConfigEntry + ) -> None: + """Initialize the cloud backup sync agent.""" + super().__init__() + self.name = config_entry.title + self._hass = hass + self._folder_id = config_entry.unique_id + self._auth: AsyncConfigEntryAuth = config_entry.runtime_data + + async def async_upload_backup( + self, + *, + open_stream: Callable[[], Coroutine[Any, Any, AsyncIterator[bytes]]], + backup: AgentBackup, + **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. + """ + properties = _convert_agent_backup_to_properties(backup) + async with Aiogoogle(user_creds=await self._creds()) as aiogoogle: + drive_v3 = await aiogoogle.discover("drive", "v3") + req = drive_v3.files.create( + pipe_from=await open_stream(), + fields="", + json={ + "name": f"{backup.name} {backup.date}.tar", + "parents": [self._folder_id], + "properties": properties, + }, + ) + try: + await aiogoogle.as_user(req, timeout=12 * 3600) + except (AiogoogleError, TimeoutError) as err: + raise BackupAgentError("Failed to upload backup") from err + + async def async_list_backups(self, **kwargs: Any) -> list[AgentBackup]: + """List backups.""" + async with Aiogoogle(user_creds=await self._creds()) as aiogoogle: + drive_v3 = await aiogoogle.discover("drive", "v3") + query = f"'{self._folder_id}' in parents and trashed=false" + try: + res = await aiogoogle.as_user( + drive_v3.files.list(q=query, fields="files(properties)"), + full_res=True, + ) + except AiogoogleError as err: + raise BackupAgentError("Failed to list backups") from err + backups = [] + async for page in res: + for file in page["files"]: + if "properties" not in file or "backup_id" not in file["properties"]: + continue + backup = _convert_properties_to_agent_backup(file["properties"]) + backups.append(backup) + return backups + + async def async_get_backup( + self, + backup_id: str, + **kwargs: Any, + ) -> AgentBackup | None: + """Return a backup.""" + async with Aiogoogle(user_creds=await self._creds()) as aiogoogle: + drive_v3 = await aiogoogle.discover("drive", "v3") + try: + res = await aiogoogle.as_user( + drive_v3.files.list( + q=self._query(backup_id), + fields="files(properties)", + ) + ) + except AiogoogleError as err: + raise BackupAgentError("Failed to get backup") from err + for file in res["files"]: + return _convert_properties_to_agent_backup(file["properties"]) + return None + + async def async_download_backup( + self, + backup_id: str, + **kwargs: Any, + ) -> AsyncIterator[bytes]: + """Download a backup file. + + :param backup_id: The ID of the backup that was returned in async_list_backups. + :return: An async iterator that yields bytes. + """ + user_creds = await self._creds() + async with Aiogoogle(user_creds=user_creds) as aiogoogle: + drive_v3 = await aiogoogle.discover("drive", "v3") + try: + res = await aiogoogle.as_user( + drive_v3.files.list( + q=self._query(backup_id), + fields="files(id)", + ) + ) + except AiogoogleError as err: + raise BackupAgentError("Failed to get backup") from err + for file in res["files"]: + # Intentionally not passing pipe_to and not wrapping this via aiogoogle.as_user + # to avoid downloading the whole file in memory + req = drive_v3.files.get(fileId=file["id"], alt="media") + req = aiogoogle.oauth2.authorize(req, user_creds) + try: + resp = await async_get_clientsession(self._hass).get( + req.url, headers=req.headers + ) + resp.raise_for_status() + except ClientError as err: + raise BackupAgentError("Failed to download backup") from err + return ChunkAsyncStreamIterator(resp.content) + raise BackupAgentError("Backup not found") + + async def async_delete_backup( + self, + backup_id: str, + **kwargs: Any, + ) -> None: + """Delete a backup file. + + :param backup_id: The ID of the backup that was returned in async_list_backups. + """ + async with Aiogoogle(user_creds=await self._creds()) as aiogoogle: + drive_v3 = await aiogoogle.discover("drive", "v3") + try: + res = await aiogoogle.as_user( + drive_v3.files.list( + q=self._query(backup_id), + fields="files(id)", + ) + ) + except AiogoogleError as err: + raise BackupAgentError("Failed to get backup") from err + for file in res["files"]: + try: + await aiogoogle.as_user(drive_v3.files.delete(fileId=file["id"])) + except AiogoogleError as err: + raise BackupAgentError("Failed to delete backup") from err + + def _query(self, backup_id: str) -> str: + return " and ".join( + [ + f"'{self._folder_id}' in parents", + f"properties has {{ key='backup_id' and value='{backup_id}' }}", + ] + ) + + async def _creds(self) -> UserCreds: + try: + await self._auth.check_and_refresh_token() + except HomeAssistantError as err: + raise BackupAgentError("Failed to authorize") from err + return convert_to_user_creds(self._auth.oauth_session.token) diff --git a/homeassistant/components/google_drive/config_flow.py b/homeassistant/components/google_drive/config_flow.py new file mode 100644 index 00000000000000..1462f8726b39b6 --- /dev/null +++ b/homeassistant/components/google_drive/config_flow.py @@ -0,0 +1,105 @@ +"""Config flow for the Google Drive integration.""" + +from __future__ import annotations + +from collections.abc import Mapping +import logging +from typing import Any + +from aiogoogle import Aiogoogle +from aiogoogle.excs import AiogoogleError + +from homeassistant.config_entries import SOURCE_REAUTH, ConfigFlowResult +from homeassistant.const import CONF_TOKEN +from homeassistant.helpers import config_entry_oauth2_flow + +from .api import convert_to_user_creds +from .const import DEFAULT_NAME, DOMAIN, DRIVE_FOLDER_URL_PREFIX, OAUTH2_SCOPES + + +class OAuth2FlowHandler( + config_entry_oauth2_flow.AbstractOAuth2FlowHandler, domain=DOMAIN +): + """Config flow to handle Google Drive OAuth2 authentication.""" + + DOMAIN = DOMAIN + + @property + def logger(self) -> logging.Logger: + """Return logger.""" + return logging.getLogger(__name__) + + @property + def extra_authorize_data(self) -> dict[str, Any]: + """Extra data that needs to be appended to the authorize url.""" + return { + "scope": " ".join(OAUTH2_SCOPES), + # Add params to ensure we get back a refresh token + "access_type": "offline", + "prompt": "consent", + } + + async def async_step_reauth( + self, entry_data: Mapping[str, Any] + ) -> ConfigFlowResult: + """Perform reauth upon an API authentication error.""" + return await self.async_step_reauth_confirm() + + async def async_step_reauth_confirm( + self, user_input: dict[str, Any] | None = None + ) -> ConfigFlowResult: + """Confirm reauth dialog.""" + if user_input is None: + return self.async_show_form(step_id="reauth_confirm") + return await self.async_step_user() + + async def async_oauth_create_entry(self, data: dict[str, Any]) -> ConfigFlowResult: + """Create an entry for the flow, or update existing entry.""" + + user_creds = convert_to_user_creds(token=data[CONF_TOKEN]) + + if self.source == SOURCE_REAUTH: + reauth_entry = self._get_reauth_entry() + async with Aiogoogle(user_creds=user_creds) as aiogoogle: + drive_v3 = await aiogoogle.discover("drive", "v3") + try: + await aiogoogle.as_user( + drive_v3.files.get(fileId=reauth_entry.unique_id, fields="") + ) + except AiogoogleError as err: + self.logger.error( + "Could not find folder '%s%s': %s", + DRIVE_FOLDER_URL_PREFIX, + reauth_entry.unique_id, + str(err), + ) + return self.async_abort(reason="get_folder_failure") + return self.async_update_reload_and_abort(reauth_entry, data=data) + + async with Aiogoogle(user_creds=user_creds) as aiogoogle: + drive_v3 = await aiogoogle.discover("drive", "v3") + req = drive_v3.files.create( + fields="id", + json={ + "name": "Home Assistant", + "mimeType": "application/vnd.google-apps.folder", + # Adding a property to be able to identify this folder + # if needed in the future. + # 1 instead of true to avoid hitting char limits + # if we ever need to add more properties. + "properties": {"ha_root": "1"}, + }, + ) + try: + res = await aiogoogle.as_user(req) + except AiogoogleError as err: + self.logger.error("Error creating folder: %s", str(err)) + return self.async_abort(reason="create_folder_failure") + folder_id = res["id"] + await self.async_set_unique_id(folder_id) + self._abort_if_unique_id_configured() + return self.async_create_entry( + title=DEFAULT_NAME, + data=data, + description_placeholders={"url": f"{DRIVE_FOLDER_URL_PREFIX}{folder_id}"}, + ) diff --git a/homeassistant/components/google_drive/const.py b/homeassistant/components/google_drive/const.py new file mode 100644 index 00000000000000..d0055eea0b214e --- /dev/null +++ b/homeassistant/components/google_drive/const.py @@ -0,0 +1,11 @@ +"""Constants for the Google Drive integration.""" + +from __future__ import annotations + +DOMAIN = "google_drive" + +DEFAULT_NAME = "Google Drive" +DRIVE_FOLDER_URL_PREFIX = "https://drive.google.com/drive/folders/" +OAUTH2_SCOPES = [ + "https://www.googleapis.com/auth/drive.file", +] diff --git a/homeassistant/components/google_drive/manifest.json b/homeassistant/components/google_drive/manifest.json new file mode 100644 index 00000000000000..f487c6bfe15da3 --- /dev/null +++ b/homeassistant/components/google_drive/manifest.json @@ -0,0 +1,14 @@ +{ + "domain": "google_drive", + "name": "Google Drive", + "after_dependencies": ["backup"], + "codeowners": ["@tronikos"], + "config_flow": true, + "dependencies": ["application_credentials"], + "documentation": "https://www.home-assistant.io/integrations/google_drive", + "integration_type": "service", + "iot_class": "cloud_polling", + "loggers": ["aiogoogle"], + "quality_scale": "bronze", + "requirements": ["aiogoogle==5.13.2"] +} diff --git a/homeassistant/components/google_drive/quality_scale.yaml b/homeassistant/components/google_drive/quality_scale.yaml new file mode 100644 index 00000000000000..699b496aec2f15 --- /dev/null +++ b/homeassistant/components/google_drive/quality_scale.yaml @@ -0,0 +1,113 @@ +rules: + # Bronze + action-setup: + status: exempt + comment: No actions. + appropriate-polling: + status: exempt + comment: No polling. + brands: done + common-modules: done + config-flow-test-coverage: todo + config-flow: done + dependency-transparency: done + docs-actions: + status: exempt + comment: No actions. + docs-high-level-description: done + docs-installation-instructions: done + docs-removal-instructions: done + entity-event-setup: + status: exempt + comment: No entities. + entity-unique-id: + status: exempt + comment: No entities. + has-entity-name: + status: exempt + comment: No entities. + runtime-data: done + test-before-configure: todo + test-before-setup: todo + unique-config-entry: done + + # Silver + action-exceptions: todo + config-entry-unloading: done + docs-configuration-parameters: + status: exempt + comment: No configuration options. + docs-installation-parameters: done + entity-unavailable: + status: exempt + comment: No entities. + integration-owner: done + log-when-unavailable: + status: exempt + comment: No entities. + parallel-updates: + status: exempt + comment: No actions and no entities. + reauthentication-flow: done + test-coverage: todo + + # Gold + devices: + status: exempt + comment: No devices. + diagnostics: + status: exempt + comment: No data to diagnose. + discovery-update-info: + status: exempt + comment: No discovery. + discovery: + status: exempt + comment: No discovery. + docs-data-update: + status: exempt + comment: No updates. + docs-examples: + status: exempt + comment: | + This integration only serves backup. + docs-known-limitations: done + docs-supported-devices: + status: exempt + comment: No devices. + docs-supported-functions: done + docs-troubleshooting: done + docs-use-cases: done + dynamic-devices: + status: exempt + comment: No devices. + entity-category: + status: exempt + comment: No entities. + entity-device-class: + status: exempt + comment: No entities. + entity-disabled-by-default: + status: exempt + comment: No entities. + entity-translations: + status: exempt + comment: No entities. + exception-translations: todo + icon-translations: + status: exempt + comment: No entities. + reconfiguration-flow: + status: exempt + comment: No configuration options. + repair-issues: + status: exempt + comment: No known user-repairable issues. + stale-devices: + status: exempt + comment: No devices. + + # Platinum + async-dependency: done + inject-websession: todo + strict-typing: todo diff --git a/homeassistant/components/google_drive/strings.json b/homeassistant/components/google_drive/strings.json new file mode 100644 index 00000000000000..c8f4aa0ddd3021 --- /dev/null +++ b/homeassistant/components/google_drive/strings.json @@ -0,0 +1,36 @@ +{ + "config": { + "step": { + "pick_implementation": { + "title": "[%key:common::config_flow::title::oauth2_pick_implementation%]" + }, + "auth": { "title": "Link Google Account" }, + "reauth_confirm": { + "title": "[%key:common::config_flow::title::reauth%]", + "description": "The Google Drive integration needs to re-authenticate your account" + } + }, + "abort": { + "already_configured": "[%key:common::config_flow::abort::already_configured_account%]", + "already_in_progress": "[%key:common::config_flow::abort::already_in_progress%]", + "cannot_connect": "[%key:common::config_flow::error::cannot_connect%]", + "timeout_connect": "[%key:common::config_flow::error::timeout_connect%]", + "oauth_error": "[%key:common::config_flow::abort::oauth2_error%]", + "missing_configuration": "[%key:common::config_flow::abort::oauth2_missing_configuration%]", + "reauth_successful": "[%key:common::config_flow::abort::reauth_successful%]", + "invalid_access_token": "[%key:common::config_flow::error::invalid_access_token%]", + "unknown": "[%key:common::config_flow::error::unknown%]", + "create_folder_failure": "Error while creating Google Drive folder, see error log for details", + "get_folder_failure": "Error while getting Google Drive folder, see error log for details", + "oauth_timeout": "[%key:common::config_flow::abort::oauth2_timeout%]", + "oauth_unauthorized": "[%key:common::config_flow::abort::oauth2_unauthorized%]", + "oauth_failed": "[%key:common::config_flow::abort::oauth2_failed%]" + }, + "create_entry": { + "default": "Successfully authenticated and created Home Assistant folder at: {url}\nFeel free to rename it in Google Drive as you wish." + } + }, + "application_credentials": { + "description": "Follow the [instructions]({more_info_url}) for [OAuth consent screen]({oauth_consent_url}) to give Home Assistant access to your Google Drive. You also need to create Application Credentials linked to your account:\n1. Go to [Credentials]({oauth_creds_url}) and select **Create Credentials**.\n1. From the drop-down list select **OAuth client ID**.\n1. Select **Web application** for the Application Type." + } +} diff --git a/homeassistant/generated/application_credentials.py b/homeassistant/generated/application_credentials.py index 6b3028826dc4f1..1f26290cf98efb 100644 --- a/homeassistant/generated/application_credentials.py +++ b/homeassistant/generated/application_credentials.py @@ -9,6 +9,7 @@ "geocaching", "google", "google_assistant_sdk", + "google_drive", "google_mail", "google_photos", "google_sheets", diff --git a/homeassistant/generated/config_flows.py b/homeassistant/generated/config_flows.py index f3e82d4d085adc..a1318d36175e57 100644 --- a/homeassistant/generated/config_flows.py +++ b/homeassistant/generated/config_flows.py @@ -230,6 +230,7 @@ "google", "google_assistant_sdk", "google_cloud", + "google_drive", "google_generative_ai_conversation", "google_mail", "google_photos", diff --git a/homeassistant/generated/integrations.json b/homeassistant/generated/integrations.json index 8343b7fde9d7da..1575926076429f 100644 --- a/homeassistant/generated/integrations.json +++ b/homeassistant/generated/integrations.json @@ -2290,6 +2290,12 @@ "iot_class": "cloud_push", "name": "Google Cloud" }, + "google_drive": { + "integration_type": "service", + "config_flow": true, + "iot_class": "cloud_polling", + "name": "Google Drive" + }, "google_generative_ai_conversation": { "integration_type": "service", "config_flow": true, diff --git a/requirements_all.txt b/requirements_all.txt index afffe77339c29c..6a97de71fdf547 100644 --- a/requirements_all.txt +++ b/requirements_all.txt @@ -254,6 +254,9 @@ aioftp==0.21.3 # homeassistant.components.github aiogithubapi==24.6.0 +# homeassistant.components.google_drive +aiogoogle==5.13.2 + # homeassistant.components.guardian aioguardian==2022.07.0 diff --git a/requirements_test_all.txt b/requirements_test_all.txt index f73d8c9cac204d..ccc6df2c5261e7 100644 --- a/requirements_test_all.txt +++ b/requirements_test_all.txt @@ -239,6 +239,9 @@ aioflo==2021.11.0 # homeassistant.components.github aiogithubapi==24.6.0 +# homeassistant.components.google_drive +aiogoogle==5.13.2 + # homeassistant.components.guardian aioguardian==2022.07.0 diff --git a/tests/components/google_drive/__init__.py b/tests/components/google_drive/__init__.py new file mode 100644 index 00000000000000..7a55f70a3d679f --- /dev/null +++ b/tests/components/google_drive/__init__.py @@ -0,0 +1 @@ +"""Tests for the Google Drive integration.""" From f6493256403ac94ac514b2f118f420001892c007 Mon Sep 17 00:00:00 2001 From: tronikos Date: Sat, 4 Jan 2025 08:14:13 +0000 Subject: [PATCH 02/23] Add test_config_flow --- .../components/google_drive/config_flow.py | 65 +-- .../components/google_drive/const.py | 1 + .../google_drive/quality_scale.yaml | 10 +- .../google_drive/test_config_flow.py | 379 ++++++++++++++++++ 4 files changed, 421 insertions(+), 34 deletions(-) create mode 100644 tests/components/google_drive/test_config_flow.py diff --git a/homeassistant/components/google_drive/config_flow.py b/homeassistant/components/google_drive/config_flow.py index 1462f8726b39b6..7be3406efb6fb4 100644 --- a/homeassistant/components/google_drive/config_flow.py +++ b/homeassistant/components/google_drive/config_flow.py @@ -6,15 +6,20 @@ import logging from typing import Any -from aiogoogle import Aiogoogle -from aiogoogle.excs import AiogoogleError +from aiohttp import ClientError from homeassistant.config_entries import SOURCE_REAUTH, ConfigFlowResult -from homeassistant.const import CONF_TOKEN +from homeassistant.const import CONF_ACCESS_TOKEN, CONF_TOKEN from homeassistant.helpers import config_entry_oauth2_flow +from homeassistant.helpers.aiohttp_client import async_get_clientsession -from .api import convert_to_user_creds -from .const import DEFAULT_NAME, DOMAIN, DRIVE_FOLDER_URL_PREFIX, OAUTH2_SCOPES +from .const import ( + DEFAULT_NAME, + DOMAIN, + DRIVE_API_FILES, + DRIVE_FOLDER_URL_PREFIX, + OAUTH2_SCOPES, +) class OAuth2FlowHandler( @@ -56,30 +61,31 @@ async def async_step_reauth_confirm( async def async_oauth_create_entry(self, data: dict[str, Any]) -> ConfigFlowResult: """Create an entry for the flow, or update existing entry.""" - user_creds = convert_to_user_creds(token=data[CONF_TOKEN]) + headers = { + "Authorization": f"Bearer {data[CONF_TOKEN][CONF_ACCESS_TOKEN]}", + } if self.source == SOURCE_REAUTH: reauth_entry = self._get_reauth_entry() - async with Aiogoogle(user_creds=user_creds) as aiogoogle: - drive_v3 = await aiogoogle.discover("drive", "v3") - try: - await aiogoogle.as_user( - drive_v3.files.get(fileId=reauth_entry.unique_id, fields="") - ) - except AiogoogleError as err: - self.logger.error( - "Could not find folder '%s%s': %s", - DRIVE_FOLDER_URL_PREFIX, - reauth_entry.unique_id, - str(err), - ) - return self.async_abort(reason="get_folder_failure") + try: + resp = await async_get_clientsession(self.hass).get( + f"{DRIVE_API_FILES}/{reauth_entry.unique_id}?fields=", + headers=headers, + ) + resp.raise_for_status() + except ClientError as err: + self.logger.error( + "Could not find folder '%s%s': %s", + DRIVE_FOLDER_URL_PREFIX, + reauth_entry.unique_id, + str(err), + ) + return self.async_abort(reason="get_folder_failure") return self.async_update_reload_and_abort(reauth_entry, data=data) - async with Aiogoogle(user_creds=user_creds) as aiogoogle: - drive_v3 = await aiogoogle.discover("drive", "v3") - req = drive_v3.files.create( - fields="id", + try: + resp = await async_get_clientsession(self.hass).post( + f"{DRIVE_API_FILES}?fields=id", json={ "name": "Home Assistant", "mimeType": "application/vnd.google-apps.folder", @@ -89,12 +95,13 @@ async def async_oauth_create_entry(self, data: dict[str, Any]) -> ConfigFlowResu # if we ever need to add more properties. "properties": {"ha_root": "1"}, }, + headers=headers, ) - try: - res = await aiogoogle.as_user(req) - except AiogoogleError as err: - self.logger.error("Error creating folder: %s", str(err)) - return self.async_abort(reason="create_folder_failure") + resp.raise_for_status() + res = await resp.json() + except ClientError as err: + self.logger.error("Error creating folder: %s", str(err)) + return self.async_abort(reason="create_folder_failure") folder_id = res["id"] await self.async_set_unique_id(folder_id) self._abort_if_unique_id_configured() diff --git a/homeassistant/components/google_drive/const.py b/homeassistant/components/google_drive/const.py index d0055eea0b214e..7aa2e49cd668cb 100644 --- a/homeassistant/components/google_drive/const.py +++ b/homeassistant/components/google_drive/const.py @@ -5,6 +5,7 @@ DOMAIN = "google_drive" DEFAULT_NAME = "Google Drive" +DRIVE_API_FILES = "https://www.googleapis.com/drive/v3/files" DRIVE_FOLDER_URL_PREFIX = "https://drive.google.com/drive/folders/" OAUTH2_SCOPES = [ "https://www.googleapis.com/auth/drive.file", diff --git a/homeassistant/components/google_drive/quality_scale.yaml b/homeassistant/components/google_drive/quality_scale.yaml index 699b496aec2f15..5dbe230f531378 100644 --- a/homeassistant/components/google_drive/quality_scale.yaml +++ b/homeassistant/components/google_drive/quality_scale.yaml @@ -8,7 +8,7 @@ rules: comment: No polling. brands: done common-modules: done - config-flow-test-coverage: todo + config-flow-test-coverage: done config-flow: done dependency-transparency: done docs-actions: @@ -27,12 +27,12 @@ rules: status: exempt comment: No entities. runtime-data: done - test-before-configure: todo - test-before-setup: todo + test-before-configure: done + test-before-setup: done unique-config-entry: done # Silver - action-exceptions: todo + action-exceptions: done config-entry-unloading: done docs-configuration-parameters: status: exempt @@ -93,7 +93,7 @@ rules: entity-translations: status: exempt comment: No entities. - exception-translations: todo + exception-translations: done icon-translations: status: exempt comment: No entities. diff --git a/tests/components/google_drive/test_config_flow.py b/tests/components/google_drive/test_config_flow.py new file mode 100644 index 00000000000000..5f2f4b97fc0e22 --- /dev/null +++ b/tests/components/google_drive/test_config_flow.py @@ -0,0 +1,379 @@ +"""Test the Google Drive config flow.""" + +from unittest.mock import patch + +from aiohttp import ClientError +import pytest + +from homeassistant import config_entries +from homeassistant.components.application_credentials import ( + ClientCredential, + async_import_client_credential, +) +from homeassistant.components.google_drive.const import DOMAIN +from homeassistant.core import HomeAssistant +from homeassistant.data_entry_flow import FlowResultType +from homeassistant.helpers import config_entry_oauth2_flow +from homeassistant.setup import async_setup_component + +from tests.common import MockConfigEntry +from tests.test_util.aiohttp import AiohttpClientMocker +from tests.typing import ClientSessionGenerator + +CLIENT_ID = "1234" +CLIENT_SECRET = "5678" +GOOGLE_AUTH_URI = "https://accounts.google.com/o/oauth2/v2/auth" +GOOGLE_TOKEN_URI = "https://oauth2.googleapis.com/token" +FOLDER_ID = "google-folder-id" +TITLE = "Google Drive" + + +@pytest.fixture +async def setup_credentials(hass: HomeAssistant) -> None: + """Fixture to setup credentials.""" + assert await async_setup_component(hass, "application_credentials", {}) + await async_import_client_credential( + hass, + DOMAIN, + ClientCredential(CLIENT_ID, CLIENT_SECRET), + ) + + +@pytest.mark.usefixtures("current_request_with_host") +async def test_full_flow( + hass: HomeAssistant, + hass_client_no_auth: ClientSessionGenerator, + aioclient_mock: AiohttpClientMocker, + setup_credentials, +) -> None: + """Check full flow.""" + result = await hass.config_entries.flow.async_init( + "google_drive", context={"source": config_entries.SOURCE_USER} + ) + state = config_entry_oauth2_flow._encode_jwt( + hass, + { + "flow_id": result["flow_id"], + "redirect_uri": "https://example.com/auth/external/callback", + }, + ) + + assert result["url"] == ( + f"{GOOGLE_AUTH_URI}?response_type=code&client_id={CLIENT_ID}" + "&redirect_uri=https://example.com/auth/external/callback" + f"&state={state}&scope=https://www.googleapis.com/auth/drive.file" + "&access_type=offline&prompt=consent" + ) + + client = await hass_client_no_auth() + resp = await client.get(f"/auth/external/callback?code=abcd&state={state}") + assert resp.status == 200 + assert resp.headers["content-type"] == "text/html; charset=utf-8" + + # Prepare API response when creating the folder + aioclient_mock.post( + "https://www.googleapis.com/drive/v3/files?fields=id", + json={"id": FOLDER_ID}, + ) + + aioclient_mock.post( + GOOGLE_TOKEN_URI, + json={ + "refresh_token": "mock-refresh-token", + "access_token": "mock-access-token", + "type": "Bearer", + "expires_in": 60, + }, + ) + + with patch( + "homeassistant.components.google_drive.async_setup_entry", return_value=True + ) as mock_setup: + result = await hass.config_entries.flow.async_configure(result["flow_id"]) + + assert len(hass.config_entries.async_entries(DOMAIN)) == 1 + assert len(mock_setup.mock_calls) == 1 + assert len(aioclient_mock.mock_calls) == 2 + assert aioclient_mock.mock_calls[1][2] == { + "name": "Home Assistant", + "mimeType": "application/vnd.google-apps.folder", + "properties": {"ha_root": "1"}, + } + assert aioclient_mock.mock_calls[1][3] == { + "Authorization": "Bearer mock-access-token" + } + + assert result.get("type") is FlowResultType.CREATE_ENTRY + assert result.get("title") == TITLE + assert "result" in result + assert result.get("result").unique_id == FOLDER_ID + assert "token" in result.get("result").data + assert result.get("result").data["token"].get("access_token") == "mock-access-token" + assert ( + result.get("result").data["token"].get("refresh_token") == "mock-refresh-token" + ) + + +@pytest.mark.usefixtures("current_request_with_host") +async def test_create_folder_error( + hass: HomeAssistant, + hass_client_no_auth: ClientSessionGenerator, + aioclient_mock: AiohttpClientMocker, + setup_credentials, +) -> None: + """Test case where creating the folder fails.""" + result = await hass.config_entries.flow.async_init( + "google_drive", context={"source": config_entries.SOURCE_USER} + ) + state = config_entry_oauth2_flow._encode_jwt( + hass, + { + "flow_id": result["flow_id"], + "redirect_uri": "https://example.com/auth/external/callback", + }, + ) + + assert result["url"] == ( + f"{GOOGLE_AUTH_URI}?response_type=code&client_id={CLIENT_ID}" + "&redirect_uri=https://example.com/auth/external/callback" + f"&state={state}&scope=https://www.googleapis.com/auth/drive.file" + "&access_type=offline&prompt=consent" + ) + + client = await hass_client_no_auth() + resp = await client.get(f"/auth/external/callback?code=abcd&state={state}") + assert resp.status == 200 + assert resp.headers["content-type"] == "text/html; charset=utf-8" + + # Prepare fake exception creating the folder + aioclient_mock.post( + "https://www.googleapis.com/drive/v3/files?fields=id", + exc=ClientError, + ) + + aioclient_mock.post( + GOOGLE_TOKEN_URI, + json={ + "refresh_token": "mock-refresh-token", + "access_token": "mock-access-token", + "type": "Bearer", + "expires_in": 60, + }, + ) + + result = await hass.config_entries.flow.async_configure(result["flow_id"]) + assert result.get("type") is FlowResultType.ABORT + assert result.get("reason") == "create_folder_failure" + + +@pytest.mark.usefixtures("current_request_with_host") +async def test_reauth( + hass: HomeAssistant, + hass_client_no_auth: ClientSessionGenerator, + aioclient_mock: AiohttpClientMocker, + setup_credentials, +) -> None: + """Test the reauthentication case updates the existing config entry.""" + + config_entry = MockConfigEntry( + domain=DOMAIN, + unique_id=FOLDER_ID, + data={ + "token": { + "access_token": "mock-access-token", + }, + }, + ) + config_entry.add_to_hass(hass) + + config_entry.async_start_reauth(hass) + await hass.async_block_till_done() + + flows = hass.config_entries.flow.async_progress() + assert len(flows) == 1 + result = flows[0] + assert result["step_id"] == "reauth_confirm" + + result = await hass.config_entries.flow.async_configure(result["flow_id"], {}) + state = config_entry_oauth2_flow._encode_jwt( + hass, + { + "flow_id": result["flow_id"], + "redirect_uri": "https://example.com/auth/external/callback", + }, + ) + assert result["url"] == ( + f"{GOOGLE_AUTH_URI}?response_type=code&client_id={CLIENT_ID}" + "&redirect_uri=https://example.com/auth/external/callback" + f"&state={state}&scope=https://www.googleapis.com/auth/drive.file" + "&access_type=offline&prompt=consent" + ) + client = await hass_client_no_auth() + resp = await client.get(f"/auth/external/callback?code=abcd&state={state}") + assert resp.status == 200 + assert resp.headers["content-type"] == "text/html; charset=utf-8" + + # Config flow will lookup existing key to make sure it still exists + aioclient_mock.get( + f"https://www.googleapis.com/drive/v3/files/{FOLDER_ID}?fields=", + json={}, + ) + + aioclient_mock.post( + GOOGLE_TOKEN_URI, + json={ + "refresh_token": "mock-refresh-token", + "access_token": "updated-access-token", + "type": "Bearer", + "expires_in": 60, + }, + ) + + with patch( + "homeassistant.components.google_drive.async_setup_entry", return_value=True + ) as mock_setup: + result = await hass.config_entries.flow.async_configure(result["flow_id"]) + await hass.async_block_till_done() + + assert len(hass.config_entries.async_entries(DOMAIN)) == 1 + assert len(mock_setup.mock_calls) == 1 + + assert result.get("type") is FlowResultType.ABORT + assert result.get("reason") == "reauth_successful" + + assert config_entry.unique_id == FOLDER_ID + assert "token" in config_entry.data + # Verify access token is refreshed + assert config_entry.data["token"].get("access_token") == "updated-access-token" + assert config_entry.data["token"].get("refresh_token") == "mock-refresh-token" + + +@pytest.mark.usefixtures("current_request_with_host") +async def test_reauth_abort( + hass: HomeAssistant, + hass_client_no_auth: ClientSessionGenerator, + aioclient_mock: AiohttpClientMocker, + setup_credentials, +) -> None: + """Test failure case during reauth.""" + + config_entry = MockConfigEntry( + domain=DOMAIN, + unique_id=FOLDER_ID, + data={ + "token": { + "access_token": "mock-access-token", + }, + }, + ) + config_entry.add_to_hass(hass) + + config_entry.async_start_reauth(hass) + await hass.async_block_till_done() + + flows = hass.config_entries.flow.async_progress() + assert len(flows) == 1 + result = flows[0] + assert result["step_id"] == "reauth_confirm" + + result = await hass.config_entries.flow.async_configure(result["flow_id"], {}) + state = config_entry_oauth2_flow._encode_jwt( + hass, + { + "flow_id": result["flow_id"], + "redirect_uri": "https://example.com/auth/external/callback", + }, + ) + assert result["url"] == ( + f"{GOOGLE_AUTH_URI}?response_type=code&client_id={CLIENT_ID}" + "&redirect_uri=https://example.com/auth/external/callback" + f"&state={state}&scope=https://www.googleapis.com/auth/drive.file" + "&access_type=offline&prompt=consent" + ) + client = await hass_client_no_auth() + resp = await client.get(f"/auth/external/callback?code=abcd&state={state}") + assert resp.status == 200 + assert resp.headers["content-type"] == "text/html; charset=utf-8" + + # Simulate failure looking up existing folder + aioclient_mock.get( + f"https://www.googleapis.com/drive/v3/files/{FOLDER_ID}?fields=", + exc=ClientError, + ) + + aioclient_mock.post( + GOOGLE_TOKEN_URI, + json={ + "refresh_token": "mock-refresh-token", + "access_token": "updated-access-token", + "type": "Bearer", + "expires_in": 60, + }, + ) + + result = await hass.config_entries.flow.async_configure(result["flow_id"]) + assert result.get("type") is FlowResultType.ABORT + assert result.get("reason") == "get_folder_failure" + + +@pytest.mark.usefixtures("current_request_with_host") +async def test_already_configured( + hass: HomeAssistant, + hass_client_no_auth: ClientSessionGenerator, + aioclient_mock: AiohttpClientMocker, + setup_credentials, +) -> None: + """Test case where config flow discovers unique id was already configured.""" + config_entry = MockConfigEntry( + domain=DOMAIN, + unique_id=FOLDER_ID, + data={ + "token": { + "access_token": "mock-access-token", + }, + }, + ) + config_entry.add_to_hass(hass) + + result = await hass.config_entries.flow.async_init( + "google_drive", context={"source": config_entries.SOURCE_USER} + ) + state = config_entry_oauth2_flow._encode_jwt( + hass, + { + "flow_id": result["flow_id"], + "redirect_uri": "https://example.com/auth/external/callback", + }, + ) + + assert result["url"] == ( + f"{GOOGLE_AUTH_URI}?response_type=code&client_id={CLIENT_ID}" + "&redirect_uri=https://example.com/auth/external/callback" + f"&state={state}&scope=https://www.googleapis.com/auth/drive.file" + "&access_type=offline&prompt=consent" + ) + + client = await hass_client_no_auth() + resp = await client.get(f"/auth/external/callback?code=abcd&state={state}") + assert resp.status == 200 + assert resp.headers["content-type"] == "text/html; charset=utf-8" + + # Prepare API response when creating the folder + aioclient_mock.post( + "https://www.googleapis.com/drive/v3/files?fields=id", + json={"id": FOLDER_ID}, + ) + + aioclient_mock.post( + GOOGLE_TOKEN_URI, + json={ + "refresh_token": "mock-refresh-token", + "access_token": "mock-access-token", + "type": "Bearer", + "expires_in": 60, + }, + ) + + result = await hass.config_entries.flow.async_configure(result["flow_id"]) + assert result.get("type") is FlowResultType.ABORT + assert result.get("reason") == "already_configured" From 4d3674c14fb397cdebf4b4b49feebbad7c04c361 Mon Sep 17 00:00:00 2001 From: tronikos Date: Sat, 4 Jan 2025 11:19:21 +0000 Subject: [PATCH 03/23] Stop using aiogoogle --- .strict-typing | 1 + homeassistant/components/google_drive/api.py | 14 +- .../components/google_drive/backup.py | 193 +++++++++--------- .../components/google_drive/config_flow.py | 7 +- .../components/google_drive/const.py | 1 + .../components/google_drive/manifest.json | 4 +- .../google_drive/quality_scale.yaml | 4 +- mypy.ini | 10 + requirements_all.txt | 3 - requirements_test_all.txt | 3 - 10 files changed, 118 insertions(+), 122 deletions(-) diff --git a/.strict-typing b/.strict-typing index 0d160982a07b5a..afa944cf095033 100644 --- a/.strict-typing +++ b/.strict-typing @@ -217,6 +217,7 @@ homeassistant.components.goalzero.* homeassistant.components.google.* homeassistant.components.google_assistant_sdk.* homeassistant.components.google_cloud.* +homeassistant.components.google_drive.* homeassistant.components.google_photos.* homeassistant.components.google_sheets.* homeassistant.components.govee_ble.* diff --git a/homeassistant/components/google_drive/api.py b/homeassistant/components/google_drive/api.py index 80f623fead67d6..200f63008dd0ba 100644 --- a/homeassistant/components/google_drive/api.py +++ b/homeassistant/components/google_drive/api.py @@ -2,10 +2,6 @@ from __future__ import annotations -from datetime import datetime -from typing import Any - -from aiogoogle.auth import UserCreds from aiohttp.client_exceptions import ClientError, ClientResponseError from google.auth.exceptions import RefreshError @@ -20,14 +16,6 @@ from homeassistant.helpers import config_entry_oauth2_flow -def convert_to_user_creds(token: dict[str, Any]) -> UserCreds: - """Convert an OAuth2Session token to UserCreds.""" - return UserCreds( - access_token=token["access_token"], - expires_at=datetime.fromtimestamp(token["expires_at"]).isoformat(), - ) - - class AsyncConfigEntryAuth: """Provide Google Drive authentication tied to an OAuth2 based config entry.""" @@ -43,7 +31,7 @@ def __init__( @property def access_token(self) -> str: """Return the access token.""" - return self.oauth_session.token[CONF_ACCESS_TOKEN] + return str(self.oauth_session.token[CONF_ACCESS_TOKEN]) async def check_and_refresh_token(self) -> str: """Check the token.""" diff --git a/homeassistant/components/google_drive/backup.py b/homeassistant/components/google_drive/backup.py index 44b7f1df9abe36..fa6ee333666a74 100644 --- a/homeassistant/components/google_drive/backup.py +++ b/homeassistant/components/google_drive/backup.py @@ -6,10 +6,7 @@ import json from typing import Any, Self -from aiogoogle import Aiogoogle -from aiogoogle.auth import UserCreds -from aiogoogle.excs import AiogoogleError -from aiohttp import ClientError, StreamReader +from aiohttp import ClientError, MultipartWriter, StreamReader from homeassistant.components.backup import AgentBackup, BackupAgent, BackupAgentError from homeassistant.core import HomeAssistant, callback @@ -17,8 +14,8 @@ from homeassistant.helpers.aiohttp_client import async_get_clientsession from . import DATA_BACKUP_AGENT_LISTENERS, GoogleDriveConfigEntry -from .api import AsyncConfigEntryAuth, convert_to_user_creds -from .const import DOMAIN +from .api import AsyncConfigEntryAuth +from .const import DOMAIN, DRIVE_API_FILES, DRIVE_API_UPLOAD_FILES # Google Drive only supports string key value pairs as properties. # We convert any other fields to JSON strings. @@ -134,42 +131,54 @@ async def async_upload_backup( :param open_stream: A function returning an async iterator that yields bytes. :param backup: Metadata about the backup that should be uploaded. """ + headers = await self._async_headers() properties = _convert_agent_backup_to_properties(backup) - async with Aiogoogle(user_creds=await self._creds()) as aiogoogle: - drive_v3 = await aiogoogle.discover("drive", "v3") - req = drive_v3.files.create( - pipe_from=await open_stream(), - fields="", - json={ + with MultipartWriter() as mpwriter: + mpwriter.append_json( + { "name": f"{backup.name} {backup.date}.tar", "parents": [self._folder_id], "properties": properties, - }, + } + ) + mpwriter.append(await open_stream()) + headers.update( + {"Content-Type": f"multipart/related; boundary={mpwriter.boundary}"} ) try: - await aiogoogle.as_user(req, timeout=12 * 3600) - except (AiogoogleError, TimeoutError) as err: + resp = await async_get_clientsession(self._hass).post( + DRIVE_API_UPLOAD_FILES, + params={"fields": ""}, + data=mpwriter, + headers=headers, + ) + resp.raise_for_status() + await resp.json() + except ClientError as err: raise BackupAgentError("Failed to upload backup") from err async def async_list_backups(self, **kwargs: Any) -> list[AgentBackup]: """List backups.""" - async with Aiogoogle(user_creds=await self._creds()) as aiogoogle: - drive_v3 = await aiogoogle.discover("drive", "v3") - query = f"'{self._folder_id}' in parents and trashed=false" - try: - res = await aiogoogle.as_user( - drive_v3.files.list(q=query, fields="files(properties)"), - full_res=True, - ) - except AiogoogleError as err: - raise BackupAgentError("Failed to list backups") from err + headers = await self._async_headers() + try: + resp = await async_get_clientsession(self._hass).get( + DRIVE_API_FILES, + params={ + "q": f"'{self._folder_id}' in parents and trashed=false", + "fields": "files(properties)", + }, + headers=headers, + ) + resp.raise_for_status() + res = await resp.json() + except ClientError as err: + raise BackupAgentError("Failed to list backups") from err backups = [] - async for page in res: - for file in page["files"]: - if "properties" not in file or "backup_id" not in file["properties"]: - continue - backup = _convert_properties_to_agent_backup(file["properties"]) - backups.append(backup) + for file in res["files"]: + if "properties" not in file or "backup_id" not in file["properties"]: + continue + backup = _convert_properties_to_agent_backup(file["properties"]) + backups.append(backup) return backups async def async_get_backup( @@ -178,20 +187,9 @@ async def async_get_backup( **kwargs: Any, ) -> AgentBackup | None: """Return a backup.""" - async with Aiogoogle(user_creds=await self._creds()) as aiogoogle: - drive_v3 = await aiogoogle.discover("drive", "v3") - try: - res = await aiogoogle.as_user( - drive_v3.files.list( - q=self._query(backup_id), - fields="files(properties)", - ) - ) - except AiogoogleError as err: - raise BackupAgentError("Failed to get backup") from err - for file in res["files"]: - return _convert_properties_to_agent_backup(file["properties"]) - return None + headers = await self._async_headers() + _, backup = await self._async_get_file_id_and_properties(backup_id, headers) + return backup async def async_download_backup( self, @@ -203,32 +201,20 @@ async def async_download_backup( :param backup_id: The ID of the backup that was returned in async_list_backups. :return: An async iterator that yields bytes. """ - user_creds = await self._creds() - async with Aiogoogle(user_creds=user_creds) as aiogoogle: - drive_v3 = await aiogoogle.discover("drive", "v3") - try: - res = await aiogoogle.as_user( - drive_v3.files.list( - q=self._query(backup_id), - fields="files(id)", - ) - ) - except AiogoogleError as err: - raise BackupAgentError("Failed to get backup") from err - for file in res["files"]: - # Intentionally not passing pipe_to and not wrapping this via aiogoogle.as_user - # to avoid downloading the whole file in memory - req = drive_v3.files.get(fileId=file["id"], alt="media") - req = aiogoogle.oauth2.authorize(req, user_creds) - try: - resp = await async_get_clientsession(self._hass).get( - req.url, headers=req.headers - ) - resp.raise_for_status() - except ClientError as err: - raise BackupAgentError("Failed to download backup") from err - return ChunkAsyncStreamIterator(resp.content) - raise BackupAgentError("Backup not found") + headers = await self._async_headers() + file_id, _ = await self._async_get_file_id_and_properties(backup_id, headers) + if file_id is None: + raise BackupAgentError("Backup not found") + try: + resp = await async_get_clientsession(self._hass).get( + f"{DRIVE_API_FILES}/{file_id}", + params={"alt": "media"}, + headers=headers, + ) + resp.raise_for_status() + except ClientError as err: + raise BackupAgentError("Failed to download backup") from err + return ChunkAsyncStreamIterator(resp.content) async def async_delete_backup( self, @@ -239,34 +225,49 @@ async def async_delete_backup( :param backup_id: The ID of the backup that was returned in async_list_backups. """ - async with Aiogoogle(user_creds=await self._creds()) as aiogoogle: - drive_v3 = await aiogoogle.discover("drive", "v3") - try: - res = await aiogoogle.as_user( - drive_v3.files.list( - q=self._query(backup_id), - fields="files(id)", - ) - ) - except AiogoogleError as err: - raise BackupAgentError("Failed to get backup") from err - for file in res["files"]: - try: - await aiogoogle.as_user(drive_v3.files.delete(fileId=file["id"])) - except AiogoogleError as err: - raise BackupAgentError("Failed to delete backup") from err - - def _query(self, backup_id: str) -> str: - return " and ".join( + headers = await self._async_headers() + file_id, _ = await self._async_get_file_id_and_properties(backup_id, headers) + if file_id is None: + return + try: + resp = await async_get_clientsession(self._hass).delete( + f"{DRIVE_API_FILES}/{file_id}", + headers=headers, + ) + resp.raise_for_status() + await resp.json() + except ClientError as err: + raise BackupAgentError("Failed to delete backup") from err + + async def _async_headers(self) -> dict[str, str]: + try: + access_token = await self._auth.check_and_refresh_token() + except HomeAssistantError as err: + raise BackupAgentError("Failed to refresh token") from err + return {"Authorization": f"Bearer {access_token}"} + + async def _async_get_file_id_and_properties( + self, backup_id: str, headers: dict[str, str] + ) -> tuple[str | None, AgentBackup | None]: + query = " and ".join( [ f"'{self._folder_id}' in parents", f"properties has {{ key='backup_id' and value='{backup_id}' }}", ] ) - - async def _creds(self) -> UserCreds: try: - await self._auth.check_and_refresh_token() - except HomeAssistantError as err: - raise BackupAgentError("Failed to authorize") from err - return convert_to_user_creds(self._auth.oauth_session.token) + resp = await async_get_clientsession(self._hass).get( + DRIVE_API_FILES, + params={ + "q": query, + "fields": "files(id,properties)", + }, + headers=headers, + ) + resp.raise_for_status() + res = await resp.json() + except ClientError as err: + raise BackupAgentError("Failed to get backup") from err + for file in res["files"]: + return file["id"], _convert_properties_to_agent_backup(file["properties"]) + return None, None diff --git a/homeassistant/components/google_drive/config_flow.py b/homeassistant/components/google_drive/config_flow.py index 7be3406efb6fb4..e5540abdfeca88 100644 --- a/homeassistant/components/google_drive/config_flow.py +++ b/homeassistant/components/google_drive/config_flow.py @@ -69,10 +69,12 @@ async def async_oauth_create_entry(self, data: dict[str, Any]) -> ConfigFlowResu reauth_entry = self._get_reauth_entry() try: resp = await async_get_clientsession(self.hass).get( - f"{DRIVE_API_FILES}/{reauth_entry.unique_id}?fields=", + f"{DRIVE_API_FILES}/{reauth_entry.unique_id}", + params={"fields": ""}, headers=headers, ) resp.raise_for_status() + await resp.json() except ClientError as err: self.logger.error( "Could not find folder '%s%s': %s", @@ -85,7 +87,8 @@ async def async_oauth_create_entry(self, data: dict[str, Any]) -> ConfigFlowResu try: resp = await async_get_clientsession(self.hass).post( - f"{DRIVE_API_FILES}?fields=id", + DRIVE_API_FILES, + params={"fields": "id"}, json={ "name": "Home Assistant", "mimeType": "application/vnd.google-apps.folder", diff --git a/homeassistant/components/google_drive/const.py b/homeassistant/components/google_drive/const.py index 7aa2e49cd668cb..39eab5b4acf9cc 100644 --- a/homeassistant/components/google_drive/const.py +++ b/homeassistant/components/google_drive/const.py @@ -6,6 +6,7 @@ DEFAULT_NAME = "Google Drive" DRIVE_API_FILES = "https://www.googleapis.com/drive/v3/files" +DRIVE_API_UPLOAD_FILES = "https://www.googleapis.com/upload/drive/v3/files" DRIVE_FOLDER_URL_PREFIX = "https://drive.google.com/drive/folders/" OAUTH2_SCOPES = [ "https://www.googleapis.com/auth/drive.file", diff --git a/homeassistant/components/google_drive/manifest.json b/homeassistant/components/google_drive/manifest.json index f487c6bfe15da3..4bff9101480131 100644 --- a/homeassistant/components/google_drive/manifest.json +++ b/homeassistant/components/google_drive/manifest.json @@ -8,7 +8,5 @@ "documentation": "https://www.home-assistant.io/integrations/google_drive", "integration_type": "service", "iot_class": "cloud_polling", - "loggers": ["aiogoogle"], - "quality_scale": "bronze", - "requirements": ["aiogoogle==5.13.2"] + "quality_scale": "bronze" } diff --git a/homeassistant/components/google_drive/quality_scale.yaml b/homeassistant/components/google_drive/quality_scale.yaml index 5dbe230f531378..64cc1fac65499c 100644 --- a/homeassistant/components/google_drive/quality_scale.yaml +++ b/homeassistant/components/google_drive/quality_scale.yaml @@ -109,5 +109,5 @@ rules: # Platinum async-dependency: done - inject-websession: todo - strict-typing: todo + inject-websession: done + strict-typing: done diff --git a/mypy.ini b/mypy.ini index 8600e5ba165eb3..b88736eaaaaf09 100644 --- a/mypy.ini +++ b/mypy.ini @@ -1926,6 +1926,16 @@ disallow_untyped_defs = true warn_return_any = true warn_unreachable = true +[mypy-homeassistant.components.google_drive.*] +check_untyped_defs = true +disallow_incomplete_defs = true +disallow_subclassing_any = true +disallow_untyped_calls = true +disallow_untyped_decorators = true +disallow_untyped_defs = true +warn_return_any = true +warn_unreachable = true + [mypy-homeassistant.components.google_photos.*] check_untyped_defs = true disallow_incomplete_defs = true diff --git a/requirements_all.txt b/requirements_all.txt index 6a97de71fdf547..afffe77339c29c 100644 --- a/requirements_all.txt +++ b/requirements_all.txt @@ -254,9 +254,6 @@ aioftp==0.21.3 # homeassistant.components.github aiogithubapi==24.6.0 -# homeassistant.components.google_drive -aiogoogle==5.13.2 - # homeassistant.components.guardian aioguardian==2022.07.0 diff --git a/requirements_test_all.txt b/requirements_test_all.txt index ccc6df2c5261e7..f73d8c9cac204d 100644 --- a/requirements_test_all.txt +++ b/requirements_test_all.txt @@ -239,9 +239,6 @@ aioflo==2021.11.0 # homeassistant.components.github aiogithubapi==24.6.0 -# homeassistant.components.google_drive -aiogoogle==5.13.2 - # homeassistant.components.guardian aioguardian==2022.07.0 From d049e9d4276e0201cfc11f30004ef765ef8c9b72 Mon Sep 17 00:00:00 2001 From: tronikos Date: Sun, 5 Jan 2025 07:09:34 +0000 Subject: [PATCH 04/23] address a few comments --- homeassistant/components/google_drive/backup.py | 5 +---- homeassistant/components/google_drive/config_flow.py | 5 +---- 2 files changed, 2 insertions(+), 8 deletions(-) diff --git a/homeassistant/components/google_drive/backup.py b/homeassistant/components/google_drive/backup.py index fa6ee333666a74..391774ab7165f0 100644 --- a/homeassistant/components/google_drive/backup.py +++ b/homeassistant/components/google_drive/backup.py @@ -14,7 +14,6 @@ from homeassistant.helpers.aiohttp_client import async_get_clientsession from . import DATA_BACKUP_AGENT_LISTENERS, GoogleDriveConfigEntry -from .api import AsyncConfigEntryAuth from .const import DOMAIN, DRIVE_API_FILES, DRIVE_API_UPLOAD_FILES # Google Drive only supports string key value pairs as properties. @@ -117,7 +116,7 @@ def __init__( self.name = config_entry.title self._hass = hass self._folder_id = config_entry.unique_id - self._auth: AsyncConfigEntryAuth = config_entry.runtime_data + self._auth = config_entry.runtime_data async def async_upload_backup( self, @@ -153,7 +152,6 @@ async def async_upload_backup( headers=headers, ) resp.raise_for_status() - await resp.json() except ClientError as err: raise BackupAgentError("Failed to upload backup") from err @@ -235,7 +233,6 @@ async def async_delete_backup( headers=headers, ) resp.raise_for_status() - await resp.json() except ClientError as err: raise BackupAgentError("Failed to delete backup") from err diff --git a/homeassistant/components/google_drive/config_flow.py b/homeassistant/components/google_drive/config_flow.py index e5540abdfeca88..24d857db29c441 100644 --- a/homeassistant/components/google_drive/config_flow.py +++ b/homeassistant/components/google_drive/config_flow.py @@ -74,7 +74,6 @@ async def async_oauth_create_entry(self, data: dict[str, Any]) -> ConfigFlowResu headers=headers, ) resp.raise_for_status() - await resp.json() except ClientError as err: self.logger.error( "Could not find folder '%s%s': %s", @@ -94,9 +93,7 @@ async def async_oauth_create_entry(self, data: dict[str, Any]) -> ConfigFlowResu "mimeType": "application/vnd.google-apps.folder", # Adding a property to be able to identify this folder # if needed in the future. - # 1 instead of true to avoid hitting char limits - # if we ever need to add more properties. - "properties": {"ha_root": "1"}, + "properties": {"ha": "root"}, }, headers=headers, ) From ce8867ec62944a824e31a2c5c31bff6c4365bd90 Mon Sep 17 00:00:00 2001 From: tronikos Date: Sun, 5 Jan 2025 09:01:56 +0000 Subject: [PATCH 05/23] Check folder exists in setup --- .../components/google_drive/__init__.py | 25 +++++++++++++++++-- homeassistant/components/google_drive/api.py | 25 +++++++++++++++++++ .../components/google_drive/backup.py | 3 ++- .../components/google_drive/config_flow.py | 18 ++++++------- .../components/google_drive/strings.json | 10 +++++++- 5 files changed, 66 insertions(+), 15 deletions(-) diff --git a/homeassistant/components/google_drive/__init__.py b/homeassistant/components/google_drive/__init__.py index b3ca0e932e6611..97428aaa1e45c8 100644 --- a/homeassistant/components/google_drive/__init__.py +++ b/homeassistant/components/google_drive/__init__.py @@ -4,15 +4,19 @@ from collections.abc import Callable +from aiohttp.client_exceptions import ClientError, ClientResponseError + from homeassistant.config_entries import ConfigEntry from homeassistant.core import HomeAssistant +from homeassistant.exceptions import ConfigEntryError, ConfigEntryNotReady +from homeassistant.helpers.aiohttp_client import async_get_clientsession from homeassistant.helpers.config_entry_oauth2_flow import ( OAuth2Session, async_get_config_entry_implementation, ) from homeassistant.util.hass_dict import HassKey -from .api import AsyncConfigEntryAuth +from .api import AsyncConfigEntryAuth, async_check_file_exists, create_headers from .const import DOMAIN DATA_BACKUP_AGENT_LISTENERS: HassKey[list[Callable[[], None]]] = HassKey( @@ -27,7 +31,24 @@ async def async_setup_entry(hass: HomeAssistant, entry: GoogleDriveConfigEntry) implementation = await async_get_config_entry_implementation(hass, entry) session = OAuth2Session(hass, entry, implementation) auth = AsyncConfigEntryAuth(hass, session) - await auth.check_and_refresh_token() + access_token = await auth.check_and_refresh_token() + try: + assert entry.unique_id + await async_check_file_exists( + async_get_clientsession(hass), + create_headers(access_token), + entry.unique_id, + ) + except ClientError as err: + if isinstance(err, ClientResponseError) and 400 <= err.status < 500: + if err.status == 404: + raise ConfigEntryError( + translation_key="config_entry_error_folder_not_found" + ) from err + raise ConfigEntryError( + translation_key="config_entry_error_folder_4xx" + ) from err + raise ConfigEntryNotReady from err entry.runtime_data = auth return True diff --git a/homeassistant/components/google_drive/api.py b/homeassistant/components/google_drive/api.py index 200f63008dd0ba..5797a0019b1be9 100644 --- a/homeassistant/components/google_drive/api.py +++ b/homeassistant/components/google_drive/api.py @@ -2,6 +2,7 @@ from __future__ import annotations +from aiohttp import ClientSession from aiohttp.client_exceptions import ClientError, ClientResponseError from google.auth.exceptions import RefreshError @@ -15,6 +16,30 @@ ) from homeassistant.helpers import config_entry_oauth2_flow +from .const import DRIVE_API_FILES + + +def create_headers(access_token: str) -> dict[str, str]: + """Create headers with the provided access token.""" + return { + "Authorization": f"Bearer {access_token}", + } + + +async def async_check_file_exists( + session: ClientSession, headers: dict[str, str], file_id: str +) -> None: + """Check the provided file or folder exists. + + :raises ClientError: if there is any error, including 404 + """ + resp = await session.get( + f"{DRIVE_API_FILES}/{file_id}", + params={"fields": ""}, + headers=headers, + ) + resp.raise_for_status() + class AsyncConfigEntryAuth: """Provide Google Drive authentication tied to an OAuth2 based config entry.""" diff --git a/homeassistant/components/google_drive/backup.py b/homeassistant/components/google_drive/backup.py index 391774ab7165f0..ee606997365e79 100644 --- a/homeassistant/components/google_drive/backup.py +++ b/homeassistant/components/google_drive/backup.py @@ -14,6 +14,7 @@ from homeassistant.helpers.aiohttp_client import async_get_clientsession from . import DATA_BACKUP_AGENT_LISTENERS, GoogleDriveConfigEntry +from .api import create_headers from .const import DOMAIN, DRIVE_API_FILES, DRIVE_API_UPLOAD_FILES # Google Drive only supports string key value pairs as properties. @@ -241,7 +242,7 @@ async def _async_headers(self) -> dict[str, str]: access_token = await self._auth.check_and_refresh_token() except HomeAssistantError as err: raise BackupAgentError("Failed to refresh token") from err - return {"Authorization": f"Bearer {access_token}"} + return create_headers(access_token) async def _async_get_file_id_and_properties( self, backup_id: str, headers: dict[str, str] diff --git a/homeassistant/components/google_drive/config_flow.py b/homeassistant/components/google_drive/config_flow.py index 24d857db29c441..678eaeeb7dcc1f 100644 --- a/homeassistant/components/google_drive/config_flow.py +++ b/homeassistant/components/google_drive/config_flow.py @@ -6,13 +6,14 @@ import logging from typing import Any -from aiohttp import ClientError +from aiohttp.client_exceptions import ClientError from homeassistant.config_entries import SOURCE_REAUTH, ConfigFlowResult from homeassistant.const import CONF_ACCESS_TOKEN, CONF_TOKEN from homeassistant.helpers import config_entry_oauth2_flow from homeassistant.helpers.aiohttp_client import async_get_clientsession +from .api import async_check_file_exists, create_headers from .const import ( DEFAULT_NAME, DOMAIN, @@ -61,19 +62,14 @@ async def async_step_reauth_confirm( async def async_oauth_create_entry(self, data: dict[str, Any]) -> ConfigFlowResult: """Create an entry for the flow, or update existing entry.""" - headers = { - "Authorization": f"Bearer {data[CONF_TOKEN][CONF_ACCESS_TOKEN]}", - } + session = async_get_clientsession(self.hass) + headers = create_headers(data[CONF_TOKEN][CONF_ACCESS_TOKEN]) if self.source == SOURCE_REAUTH: reauth_entry = self._get_reauth_entry() + assert reauth_entry.unique_id try: - resp = await async_get_clientsession(self.hass).get( - f"{DRIVE_API_FILES}/{reauth_entry.unique_id}", - params={"fields": ""}, - headers=headers, - ) - resp.raise_for_status() + await async_check_file_exists(session, headers, reauth_entry.unique_id) except ClientError as err: self.logger.error( "Could not find folder '%s%s': %s", @@ -85,7 +81,7 @@ async def async_oauth_create_entry(self, data: dict[str, Any]) -> ConfigFlowResu return self.async_update_reload_and_abort(reauth_entry, data=data) try: - resp = await async_get_clientsession(self.hass).post( + resp = await session.post( DRIVE_API_FILES, params={"fields": "id"}, json={ diff --git a/homeassistant/components/google_drive/strings.json b/homeassistant/components/google_drive/strings.json index c8f4aa0ddd3021..8cddcd17a399cc 100644 --- a/homeassistant/components/google_drive/strings.json +++ b/homeassistant/components/google_drive/strings.json @@ -27,10 +27,18 @@ "oauth_failed": "[%key:common::config_flow::abort::oauth2_failed%]" }, "create_entry": { - "default": "Successfully authenticated and created Home Assistant folder at: {url}\nFeel free to rename it in Google Drive as you wish." + "default": "Successfully authenticated and created 'Home Assistant' folder at: {url}\nFeel free to rename it in Google Drive as you wish." } }, "application_credentials": { "description": "Follow the [instructions]({more_info_url}) for [OAuth consent screen]({oauth_consent_url}) to give Home Assistant access to your Google Drive. You also need to create Application Credentials linked to your account:\n1. Go to [Credentials]({oauth_creds_url}) and select **Create Credentials**.\n1. From the drop-down list select **OAuth client ID**.\n1. Select **Web application** for the Application Type." + }, + "exceptions": { + "config_entry_error_folder_not_found": { + "message": "Google Drive folder not found. Remove the Google Drive integration and re-add it for a new 'Home Assistant' folder to be created in Google Drive." + }, + "config_entry_error_folder_4xx": { + "message": "Unexpected error. Open an issue." + } } } From 19d0975731907af9c1f1bcd28fb8a8b941ef100e Mon Sep 17 00:00:00 2001 From: tronikos Date: Sun, 5 Jan 2025 09:03:24 +0000 Subject: [PATCH 06/23] fix test --- tests/components/google_drive/test_config_flow.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/components/google_drive/test_config_flow.py b/tests/components/google_drive/test_config_flow.py index 5f2f4b97fc0e22..d1a76c7831d27f 100644 --- a/tests/components/google_drive/test_config_flow.py +++ b/tests/components/google_drive/test_config_flow.py @@ -97,7 +97,7 @@ async def test_full_flow( assert aioclient_mock.mock_calls[1][2] == { "name": "Home Assistant", "mimeType": "application/vnd.google-apps.folder", - "properties": {"ha_root": "1"}, + "properties": {"ha": "root"}, } assert aioclient_mock.mock_calls[1][3] == { "Authorization": "Bearer mock-access-token" From a858a81316b03e4cc47cf7fbab73d1fd952cfc48 Mon Sep 17 00:00:00 2001 From: tronikos Date: Sun, 5 Jan 2025 09:48:46 +0000 Subject: [PATCH 07/23] address comments --- .../components/google_drive/backup.py | 29 ++++++------------- 1 file changed, 9 insertions(+), 20 deletions(-) diff --git a/homeassistant/components/google_drive/backup.py b/homeassistant/components/google_drive/backup.py index ee606997365e79..fa3c2812ddd57a 100644 --- a/homeassistant/components/google_drive/backup.py +++ b/homeassistant/components/google_drive/backup.py @@ -6,7 +6,7 @@ import json from typing import Any, Self -from aiohttp import ClientError, MultipartWriter, StreamReader +from aiohttp import ClientError, ClientTimeout, MultipartWriter, StreamReader from homeassistant.components.backup import AgentBackup, BackupAgent, BackupAgentError from homeassistant.core import HomeAssistant, callback @@ -17,33 +17,20 @@ from .api import create_headers from .const import DOMAIN, DRIVE_API_FILES, DRIVE_API_UPLOAD_FILES -# Google Drive only supports string key value pairs as properties. -# We convert any other fields to JSON strings. -_AGENT_BACKUP_SIMPLE_FIELDS = [ - "backup_id", - "date", - "database_included", - "homeassistant_included", - "homeassistant_version", - "name", - "protected", - "size", -] +_UPLOAD_TIMEOUT = 12 * 3600 +# Google Drive only supports string key value pairs as properties. +# Convert every field to JSON strings except backup_id so that we can query it. def _convert_agent_backup_to_properties(backup: AgentBackup) -> dict[str, str]: return { - k: v if k in _AGENT_BACKUP_SIMPLE_FIELDS else json.dumps(v) - for k, v in backup.as_dict().items() + k: v if k == "backup_id" else json.dumps(v) for k, v in backup.as_dict().items() } def _convert_properties_to_agent_backup(d: dict[str, str]) -> AgentBackup: return AgentBackup.from_dict( - { - k: v if k in _AGENT_BACKUP_SIMPLE_FIELDS else json.loads(v) - for k, v in d.items() - } + {k: v if k == "backup_id" else json.loads(v) for k, v in d.items()} ) @@ -151,9 +138,11 @@ async def async_upload_backup( params={"fields": ""}, data=mpwriter, headers=headers, + timeout=ClientTimeout(total=_UPLOAD_TIMEOUT), ) resp.raise_for_status() - except ClientError as err: + await resp.json() + except (ClientError, TimeoutError) as err: raise BackupAgentError("Failed to upload backup") from err async def async_list_backups(self, **kwargs: Any) -> list[AgentBackup]: From 7778625101990fc2f912b26056c00c42fe9beb81 Mon Sep 17 00:00:00 2001 From: tronikos Date: Sun, 5 Jan 2025 09:49:51 +0000 Subject: [PATCH 08/23] fix --- homeassistant/components/google_drive/backup.py | 1 - 1 file changed, 1 deletion(-) diff --git a/homeassistant/components/google_drive/backup.py b/homeassistant/components/google_drive/backup.py index fa3c2812ddd57a..44e4ff8b312c56 100644 --- a/homeassistant/components/google_drive/backup.py +++ b/homeassistant/components/google_drive/backup.py @@ -141,7 +141,6 @@ async def async_upload_backup( timeout=ClientTimeout(total=_UPLOAD_TIMEOUT), ) resp.raise_for_status() - await resp.json() except (ClientError, TimeoutError) as err: raise BackupAgentError("Failed to upload backup") from err From 4d40a336ae980d41991becb85ea638c1a3b0c581 Mon Sep 17 00:00:00 2001 From: tronikos Date: Sun, 5 Jan 2025 10:00:38 +0000 Subject: [PATCH 09/23] fix --- tests/components/google_drive/test_config_flow.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/components/google_drive/test_config_flow.py b/tests/components/google_drive/test_config_flow.py index d1a76c7831d27f..683ecac8b90ed3 100644 --- a/tests/components/google_drive/test_config_flow.py +++ b/tests/components/google_drive/test_config_flow.py @@ -213,7 +213,7 @@ async def test_reauth( assert resp.status == 200 assert resp.headers["content-type"] == "text/html; charset=utf-8" - # Config flow will lookup existing key to make sure it still exists + # Config flow will lookup existing folder to make sure it still exists aioclient_mock.get( f"https://www.googleapis.com/drive/v3/files/{FOLDER_ID}?fields=", json={}, From 9ab464bc5c8477ddd9bc27528b9bb3eddba28dd8 Mon Sep 17 00:00:00 2001 From: tronikos Date: Mon, 6 Jan 2025 05:41:39 +0000 Subject: [PATCH 10/23] Use ChunkAsyncStreamIterator in helpers --- .../components/google_drive/backup.py | 34 ++++--------------- 1 file changed, 6 insertions(+), 28 deletions(-) diff --git a/homeassistant/components/google_drive/backup.py b/homeassistant/components/google_drive/backup.py index 44e4ff8b312c56..0fd5de4d25e832 100644 --- a/homeassistant/components/google_drive/backup.py +++ b/homeassistant/components/google_drive/backup.py @@ -4,14 +4,17 @@ from collections.abc import AsyncIterator, Callable, Coroutine import json -from typing import Any, Self +from typing import Any -from aiohttp import ClientError, ClientTimeout, MultipartWriter, StreamReader +from aiohttp import ClientError, ClientTimeout, MultipartWriter from homeassistant.components.backup import AgentBackup, BackupAgent, BackupAgentError from homeassistant.core import HomeAssistant, callback from homeassistant.exceptions import HomeAssistantError -from homeassistant.helpers.aiohttp_client import async_get_clientsession +from homeassistant.helpers.aiohttp_client import ( + ChunkAsyncStreamIterator, + async_get_clientsession, +) from . import DATA_BACKUP_AGENT_LISTENERS, GoogleDriveConfigEntry from .api import create_headers @@ -66,31 +69,6 @@ def remove_listener() -> None: return remove_listener -class ChunkAsyncStreamIterator: - """Async iterator for chunked streams. - - Based on aiohttp.streams.ChunkTupleAsyncStreamIterator, but yields - bytes instead of tuple[bytes, bool]. - """ - - __slots__ = ("_stream",) - - def __init__(self, stream: StreamReader) -> None: - """Initialize.""" - self._stream = stream - - def __aiter__(self) -> Self: - """Iterate.""" - return self - - async def __anext__(self) -> bytes: - """Yield next chunk.""" - rv = await self._stream.readchunk() - if rv == (b"", False): - raise StopAsyncIteration - return rv[0] - - class GoogleDriveBackupAgent(BackupAgent): """Google Drive backup agent.""" From eb3a3f169b5ccaf2d19dc0fba9cd6e92f571022f Mon Sep 17 00:00:00 2001 From: tronikos Date: Mon, 6 Jan 2025 05:42:35 +0000 Subject: [PATCH 11/23] repair-issues: todo --- homeassistant/components/google_drive/quality_scale.yaml | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/homeassistant/components/google_drive/quality_scale.yaml b/homeassistant/components/google_drive/quality_scale.yaml index 64cc1fac65499c..7f7b7bd76e1462 100644 --- a/homeassistant/components/google_drive/quality_scale.yaml +++ b/homeassistant/components/google_drive/quality_scale.yaml @@ -100,9 +100,7 @@ rules: reconfiguration-flow: status: exempt comment: No configuration options. - repair-issues: - status: exempt - comment: No known user-repairable issues. + repair-issues: todo stale-devices: status: exempt comment: No devices. From 9fa216cba10a288700d83848513cd2ccdcb614d7 Mon Sep 17 00:00:00 2001 From: tronikos Date: Mon, 6 Jan 2025 21:54:56 +0000 Subject: [PATCH 12/23] Remove check if folder exists in the reatuh flow. This is done in setup. --- .../components/google_drive/__init__.py | 14 ++-- homeassistant/components/google_drive/api.py | 18 ----- .../components/google_drive/config_flow.py | 18 +---- .../components/google_drive/strings.json | 1 - .../google_drive/test_config_flow.py | 74 ------------------- 5 files changed, 11 insertions(+), 114 deletions(-) diff --git a/homeassistant/components/google_drive/__init__.py b/homeassistant/components/google_drive/__init__.py index 97428aaa1e45c8..729bb851b9c022 100644 --- a/homeassistant/components/google_drive/__init__.py +++ b/homeassistant/components/google_drive/__init__.py @@ -16,8 +16,8 @@ ) from homeassistant.util.hass_dict import HassKey -from .api import AsyncConfigEntryAuth, async_check_file_exists, create_headers -from .const import DOMAIN +from .api import AsyncConfigEntryAuth, create_headers +from .const import DOMAIN, DRIVE_API_FILES DATA_BACKUP_AGENT_LISTENERS: HassKey[list[Callable[[], None]]] = HassKey( f"{DOMAIN}.backup_agent_listeners" @@ -33,12 +33,12 @@ async def async_setup_entry(hass: HomeAssistant, entry: GoogleDriveConfigEntry) auth = AsyncConfigEntryAuth(hass, session) access_token = await auth.check_and_refresh_token() try: - assert entry.unique_id - await async_check_file_exists( - async_get_clientsession(hass), - create_headers(access_token), - entry.unique_id, + resp = await async_get_clientsession(hass).get( + f"{DRIVE_API_FILES}/{entry.unique_id}", + params={"fields": ""}, + headers=create_headers(access_token), ) + resp.raise_for_status() except ClientError as err: if isinstance(err, ClientResponseError) and 400 <= err.status < 500: if err.status == 404: diff --git a/homeassistant/components/google_drive/api.py b/homeassistant/components/google_drive/api.py index 5797a0019b1be9..5328e3c2793300 100644 --- a/homeassistant/components/google_drive/api.py +++ b/homeassistant/components/google_drive/api.py @@ -2,7 +2,6 @@ from __future__ import annotations -from aiohttp import ClientSession from aiohttp.client_exceptions import ClientError, ClientResponseError from google.auth.exceptions import RefreshError @@ -16,8 +15,6 @@ ) from homeassistant.helpers import config_entry_oauth2_flow -from .const import DRIVE_API_FILES - def create_headers(access_token: str) -> dict[str, str]: """Create headers with the provided access token.""" @@ -26,21 +23,6 @@ def create_headers(access_token: str) -> dict[str, str]: } -async def async_check_file_exists( - session: ClientSession, headers: dict[str, str], file_id: str -) -> None: - """Check the provided file or folder exists. - - :raises ClientError: if there is any error, including 404 - """ - resp = await session.get( - f"{DRIVE_API_FILES}/{file_id}", - params={"fields": ""}, - headers=headers, - ) - resp.raise_for_status() - - class AsyncConfigEntryAuth: """Provide Google Drive authentication tied to an OAuth2 based config entry.""" diff --git a/homeassistant/components/google_drive/config_flow.py b/homeassistant/components/google_drive/config_flow.py index 678eaeeb7dcc1f..6a794dc3ae3eda 100644 --- a/homeassistant/components/google_drive/config_flow.py +++ b/homeassistant/components/google_drive/config_flow.py @@ -13,7 +13,7 @@ from homeassistant.helpers import config_entry_oauth2_flow from homeassistant.helpers.aiohttp_client import async_get_clientsession -from .api import async_check_file_exists, create_headers +from .api import create_headers from .const import ( DEFAULT_NAME, DOMAIN, @@ -66,19 +66,9 @@ async def async_oauth_create_entry(self, data: dict[str, Any]) -> ConfigFlowResu headers = create_headers(data[CONF_TOKEN][CONF_ACCESS_TOKEN]) if self.source == SOURCE_REAUTH: - reauth_entry = self._get_reauth_entry() - assert reauth_entry.unique_id - try: - await async_check_file_exists(session, headers, reauth_entry.unique_id) - except ClientError as err: - self.logger.error( - "Could not find folder '%s%s': %s", - DRIVE_FOLDER_URL_PREFIX, - reauth_entry.unique_id, - str(err), - ) - return self.async_abort(reason="get_folder_failure") - return self.async_update_reload_and_abort(reauth_entry, data=data) + return self.async_update_reload_and_abort( + self._get_reauth_entry(), data=data + ) try: resp = await session.post( diff --git a/homeassistant/components/google_drive/strings.json b/homeassistant/components/google_drive/strings.json index 8cddcd17a399cc..f6d14219d6c864 100644 --- a/homeassistant/components/google_drive/strings.json +++ b/homeassistant/components/google_drive/strings.json @@ -21,7 +21,6 @@ "invalid_access_token": "[%key:common::config_flow::error::invalid_access_token%]", "unknown": "[%key:common::config_flow::error::unknown%]", "create_folder_failure": "Error while creating Google Drive folder, see error log for details", - "get_folder_failure": "Error while getting Google Drive folder, see error log for details", "oauth_timeout": "[%key:common::config_flow::abort::oauth2_timeout%]", "oauth_unauthorized": "[%key:common::config_flow::abort::oauth2_unauthorized%]", "oauth_failed": "[%key:common::config_flow::abort::oauth2_failed%]" diff --git a/tests/components/google_drive/test_config_flow.py b/tests/components/google_drive/test_config_flow.py index 683ecac8b90ed3..6409f554c674b3 100644 --- a/tests/components/google_drive/test_config_flow.py +++ b/tests/components/google_drive/test_config_flow.py @@ -213,12 +213,6 @@ async def test_reauth( assert resp.status == 200 assert resp.headers["content-type"] == "text/html; charset=utf-8" - # Config flow will lookup existing folder to make sure it still exists - aioclient_mock.get( - f"https://www.googleapis.com/drive/v3/files/{FOLDER_ID}?fields=", - json={}, - ) - aioclient_mock.post( GOOGLE_TOKEN_URI, json={ @@ -248,74 +242,6 @@ async def test_reauth( assert config_entry.data["token"].get("refresh_token") == "mock-refresh-token" -@pytest.mark.usefixtures("current_request_with_host") -async def test_reauth_abort( - hass: HomeAssistant, - hass_client_no_auth: ClientSessionGenerator, - aioclient_mock: AiohttpClientMocker, - setup_credentials, -) -> None: - """Test failure case during reauth.""" - - config_entry = MockConfigEntry( - domain=DOMAIN, - unique_id=FOLDER_ID, - data={ - "token": { - "access_token": "mock-access-token", - }, - }, - ) - config_entry.add_to_hass(hass) - - config_entry.async_start_reauth(hass) - await hass.async_block_till_done() - - flows = hass.config_entries.flow.async_progress() - assert len(flows) == 1 - result = flows[0] - assert result["step_id"] == "reauth_confirm" - - result = await hass.config_entries.flow.async_configure(result["flow_id"], {}) - state = config_entry_oauth2_flow._encode_jwt( - hass, - { - "flow_id": result["flow_id"], - "redirect_uri": "https://example.com/auth/external/callback", - }, - ) - assert result["url"] == ( - f"{GOOGLE_AUTH_URI}?response_type=code&client_id={CLIENT_ID}" - "&redirect_uri=https://example.com/auth/external/callback" - f"&state={state}&scope=https://www.googleapis.com/auth/drive.file" - "&access_type=offline&prompt=consent" - ) - client = await hass_client_no_auth() - resp = await client.get(f"/auth/external/callback?code=abcd&state={state}") - assert resp.status == 200 - assert resp.headers["content-type"] == "text/html; charset=utf-8" - - # Simulate failure looking up existing folder - aioclient_mock.get( - f"https://www.googleapis.com/drive/v3/files/{FOLDER_ID}?fields=", - exc=ClientError, - ) - - aioclient_mock.post( - GOOGLE_TOKEN_URI, - json={ - "refresh_token": "mock-refresh-token", - "access_token": "updated-access-token", - "type": "Bearer", - "expires_in": 60, - }, - ) - - result = await hass.config_entries.flow.async_configure(result["flow_id"]) - assert result.get("type") is FlowResultType.ABORT - assert result.get("reason") == "get_folder_failure" - - @pytest.mark.usefixtures("current_request_with_host") async def test_already_configured( hass: HomeAssistant, From 3e3f106e875ca70acd72500ef45179ce6ec5601d Mon Sep 17 00:00:00 2001 From: tronikos Date: Mon, 6 Jan 2025 21:59:16 +0000 Subject: [PATCH 13/23] single_config_entry": true --- .../components/google_drive/manifest.json | 3 +- .../google_drive/test_config_flow.py | 41 +------------------ 2 files changed, 4 insertions(+), 40 deletions(-) diff --git a/homeassistant/components/google_drive/manifest.json b/homeassistant/components/google_drive/manifest.json index 4bff9101480131..32231b16616a88 100644 --- a/homeassistant/components/google_drive/manifest.json +++ b/homeassistant/components/google_drive/manifest.json @@ -8,5 +8,6 @@ "documentation": "https://www.home-assistant.io/integrations/google_drive", "integration_type": "service", "iot_class": "cloud_polling", - "quality_scale": "bronze" + "quality_scale": "bronze", + "single_config_entry": true } diff --git a/tests/components/google_drive/test_config_flow.py b/tests/components/google_drive/test_config_flow.py index 6409f554c674b3..6f51517deddce3 100644 --- a/tests/components/google_drive/test_config_flow.py +++ b/tests/components/google_drive/test_config_flow.py @@ -249,7 +249,7 @@ async def test_already_configured( aioclient_mock: AiohttpClientMocker, setup_credentials, ) -> None: - """Test case where config flow discovers unique id was already configured.""" + """Test case for single_instance_allowed.""" config_entry = MockConfigEntry( domain=DOMAIN, unique_id=FOLDER_ID, @@ -264,42 +264,5 @@ async def test_already_configured( result = await hass.config_entries.flow.async_init( "google_drive", context={"source": config_entries.SOURCE_USER} ) - state = config_entry_oauth2_flow._encode_jwt( - hass, - { - "flow_id": result["flow_id"], - "redirect_uri": "https://example.com/auth/external/callback", - }, - ) - - assert result["url"] == ( - f"{GOOGLE_AUTH_URI}?response_type=code&client_id={CLIENT_ID}" - "&redirect_uri=https://example.com/auth/external/callback" - f"&state={state}&scope=https://www.googleapis.com/auth/drive.file" - "&access_type=offline&prompt=consent" - ) - - client = await hass_client_no_auth() - resp = await client.get(f"/auth/external/callback?code=abcd&state={state}") - assert resp.status == 200 - assert resp.headers["content-type"] == "text/html; charset=utf-8" - - # Prepare API response when creating the folder - aioclient_mock.post( - "https://www.googleapis.com/drive/v3/files?fields=id", - json={"id": FOLDER_ID}, - ) - - aioclient_mock.post( - GOOGLE_TOKEN_URI, - json={ - "refresh_token": "mock-refresh-token", - "access_token": "mock-access-token", - "type": "Bearer", - "expires_in": 60, - }, - ) - - result = await hass.config_entries.flow.async_configure(result["flow_id"]) assert result.get("type") is FlowResultType.ABORT - assert result.get("reason") == "already_configured" + assert result.get("reason") == "single_instance_allowed" From ace8e577014a11c5e7ad198697eb6fcdde6498fd Mon Sep 17 00:00:00 2001 From: tronikos Date: Mon, 6 Jan 2025 23:48:39 +0000 Subject: [PATCH 14/23] Add test_init.py --- tests/components/google_drive/test_init.py | 204 +++++++++++++++++++++ 1 file changed, 204 insertions(+) create mode 100644 tests/components/google_drive/test_init.py diff --git a/tests/components/google_drive/test_init.py b/tests/components/google_drive/test_init.py new file mode 100644 index 00000000000000..714808a899c0bc --- /dev/null +++ b/tests/components/google_drive/test_init.py @@ -0,0 +1,204 @@ +"""Tests for Google Drive.""" + +from collections.abc import Awaitable, Callable, Coroutine +import http +import time +from typing import Any + +import pytest + +from homeassistant.components.application_credentials import ( + ClientCredential, + async_import_client_credential, +) +from homeassistant.components.google_drive import DOMAIN +from homeassistant.config_entries import ConfigEntryState +from homeassistant.core import HomeAssistant +from homeassistant.setup import async_setup_component + +from tests.common import MockConfigEntry +from tests.test_util.aiohttp import AiohttpClientMocker + +TEST_FOLDER_ID = "google-folder-it" + +type ComponentSetup = Callable[[], Awaitable[None]] + + +@pytest.fixture(name="expires_at") +def mock_expires_at() -> int: + """Fixture to set the oauth token expiration time.""" + return time.time() + 3600 + + +@pytest.fixture(name="config_entry") +def mock_config_entry(expires_at: int) -> MockConfigEntry: + """Fixture for MockConfigEntry.""" + return MockConfigEntry( + domain=DOMAIN, + unique_id=TEST_FOLDER_ID, + data={ + "auth_implementation": DOMAIN, + "token": { + "access_token": "mock-access-token", + "refresh_token": "mock-refresh-token", + "expires_at": expires_at, + "scope": "https://www.googleapis.com/auth/drive.file", + }, + }, + ) + + +@pytest.fixture(name="setup_integration") +async def mock_setup_integration( + hass: HomeAssistant, config_entry: MockConfigEntry +) -> Callable[[], Coroutine[Any, Any, None]]: + """Fixture for setting up the component.""" + config_entry.add_to_hass(hass) + + assert await async_setup_component(hass, "application_credentials", {}) + await async_import_client_credential( + hass, + DOMAIN, + ClientCredential("client-id", "client-secret"), + DOMAIN, + ) + + async def func() -> None: + assert await async_setup_component(hass, DOMAIN, {}) + await hass.async_block_till_done() + + return func + + +async def test_setup_success( + hass: HomeAssistant, + setup_integration: ComponentSetup, + aioclient_mock: AiohttpClientMocker, +) -> None: + """Test successful setup and unload.""" + # Setup looks up existing folder to make sure it still exists + aioclient_mock.get( + f"https://www.googleapis.com/drive/v3/files/{TEST_FOLDER_ID}?fields=", + json={}, + ) + + await setup_integration() + + entries = hass.config_entries.async_entries(DOMAIN) + assert len(entries) == 1 + assert entries[0].state is ConfigEntryState.LOADED + + await hass.config_entries.async_unload(entries[0].entry_id) + await hass.async_block_till_done() + + assert not hass.data.get(DOMAIN) + assert entries[0].state is ConfigEntryState.NOT_LOADED + assert not hass.services.async_services().get(DOMAIN, {}) + + +@pytest.mark.parametrize( + ("status", "expected_state"), + [ + ( + http.HTTPStatus.BAD_REQUEST, + ConfigEntryState.SETUP_ERROR, + ), + ( + http.HTTPStatus.NOT_FOUND, + ConfigEntryState.SETUP_ERROR, + ), + ( + http.HTTPStatus.INTERNAL_SERVER_ERROR, + ConfigEntryState.SETUP_RETRY, + ), + ], + ids=["bad_request", "not_found", "transient_failure"], +) +async def test_setup_error( + hass: HomeAssistant, + setup_integration: ComponentSetup, + aioclient_mock: AiohttpClientMocker, + status: http.HTTPStatus, + expected_state: ConfigEntryState, +) -> None: + """Test setup error.""" + # Simulate failure looking up existing folder + aioclient_mock.get( + f"https://www.googleapis.com/drive/v3/files/{TEST_FOLDER_ID}?fields=", + status=status, + ) + + await setup_integration() + + entries = hass.config_entries.async_entries(DOMAIN) + assert len(entries) == 1 + assert entries[0].state is expected_state + + +@pytest.mark.parametrize("expires_at", [time.time() - 3600], ids=["expired"]) +async def test_expired_token_refresh_success( + hass: HomeAssistant, + setup_integration: ComponentSetup, + aioclient_mock: AiohttpClientMocker, +) -> None: + """Test expired token is refreshed.""" + # Setup looks up existing folder to make sure it still exists + aioclient_mock.get( + f"https://www.googleapis.com/drive/v3/files/{TEST_FOLDER_ID}?fields=", + json={}, + ) + + aioclient_mock.post( + "https://oauth2.googleapis.com/token", + json={ + "access_token": "updated-access-token", + "refresh_token": "updated-refresh-token", + "expires_at": time.time() + 3600, + "expires_in": 3600, + }, + ) + + await setup_integration() + + entries = hass.config_entries.async_entries(DOMAIN) + assert len(entries) == 1 + assert entries[0].state is ConfigEntryState.LOADED + assert entries[0].data["token"]["access_token"] == "updated-access-token" + assert entries[0].data["token"]["expires_in"] == 3600 + + +@pytest.mark.parametrize( + ("expires_at", "status", "expected_state"), + [ + ( + time.time() - 3600, + http.HTTPStatus.UNAUTHORIZED, + ConfigEntryState.SETUP_ERROR, + ), + ( + time.time() - 3600, + http.HTTPStatus.INTERNAL_SERVER_ERROR, + ConfigEntryState.SETUP_RETRY, + ), + ], + ids=["failure_requires_reauth", "transient_failure"], +) +async def test_expired_token_refresh_failure( + hass: HomeAssistant, + setup_integration: ComponentSetup, + aioclient_mock: AiohttpClientMocker, + status: http.HTTPStatus, + expected_state: ConfigEntryState, +) -> None: + """Test failure while refreshing token with a transient error.""" + + aioclient_mock.post( + "https://oauth2.googleapis.com/token", + status=status, + ) + + await setup_integration() + + # Verify a transient failure has occurred + entries = hass.config_entries.async_entries(DOMAIN) + assert entries[0].state is expected_state From 9e6153a54fdd93e0cc66731eaa47e422b1f59889 Mon Sep 17 00:00:00 2001 From: tronikos Date: Wed, 8 Jan 2025 07:20:37 +0000 Subject: [PATCH 15/23] Store into backups.json to avoid 124 bytes per property limit --- .../components/google_drive/backup.py | 342 ++++++++++++------ .../components/google_drive/config_flow.py | 6 +- 2 files changed, 244 insertions(+), 104 deletions(-) diff --git a/homeassistant/components/google_drive/backup.py b/homeassistant/components/google_drive/backup.py index 0fd5de4d25e832..eff067c14abeb0 100644 --- a/homeassistant/components/google_drive/backup.py +++ b/homeassistant/components/google_drive/backup.py @@ -2,11 +2,13 @@ from __future__ import annotations -from collections.abc import AsyncIterator, Callable, Coroutine +from asyncio import Lock +from collections.abc import AsyncIterator, Awaitable, Callable, Coroutine import json +import logging from typing import Any -from aiohttp import ClientError, ClientTimeout, MultipartWriter +from aiohttp import ClientError, ClientTimeout, MultipartWriter, StreamReader from homeassistant.components.backup import AgentBackup, BackupAgent, BackupAgentError from homeassistant.core import HomeAssistant, callback @@ -18,34 +20,24 @@ from . import DATA_BACKUP_AGENT_LISTENERS, GoogleDriveConfigEntry from .api import create_headers -from .const import DOMAIN, DRIVE_API_FILES, DRIVE_API_UPLOAD_FILES +from .const import ( + DOMAIN, + DRIVE_API_FILES, + DRIVE_API_UPLOAD_FILES, + DRIVE_FOLDER_URL_PREFIX, +) +_LOGGER = logging.getLogger(__name__) _UPLOAD_TIMEOUT = 12 * 3600 -# Google Drive only supports string key value pairs as properties. -# Convert every field to JSON strings except backup_id so that we can query it. -def _convert_agent_backup_to_properties(backup: AgentBackup) -> dict[str, str]: - return { - k: v if k == "backup_id" else json.dumps(v) for k, v in backup.as_dict().items() - } - - -def _convert_properties_to_agent_backup(d: dict[str, str]) -> AgentBackup: - return AgentBackup.from_dict( - {k: v if k == "backup_id" else json.loads(v) for k, v in d.items()} - ) - - async def async_get_backup_agents( hass: HomeAssistant, **kwargs: Any, ) -> list[BackupAgent]: """Return a list of backup agents.""" - return [ - GoogleDriveBackupAgent(hass=hass, config_entry=config_entry) - for config_entry in hass.config_entries.async_loaded_entries(DOMAIN) - ] + entries = hass.config_entries.async_loaded_entries(DOMAIN) + return [GoogleDriveBackupAgent(hass, entry) for entry in entries] @callback @@ -83,6 +75,7 @@ def __init__( self._hass = hass self._folder_id = config_entry.unique_id self._auth = config_entry.runtime_data + self._update_backups_json_lock = Lock() async def async_upload_backup( self, @@ -97,54 +90,44 @@ async def async_upload_backup( :param backup: Metadata about the backup that should be uploaded. """ headers = await self._async_headers() - properties = _convert_agent_backup_to_properties(backup) - with MultipartWriter() as mpwriter: - mpwriter.append_json( - { - "name": f"{backup.name} {backup.date}.tar", - "parents": [self._folder_id], - "properties": properties, - } + backup_metadata = { + "name": f"{backup.name} {backup.date}.tar", + "parents": [self._folder_id], + "properties": { + "ha": "backup", + "backup_id": backup.backup_id, + }, + } + try: + _LOGGER.debug( + "Uploading backup: %s with Google Drive metadata: %s", + backup, + backup_metadata, ) - mpwriter.append(await open_stream()) - headers.update( - {"Content-Type": f"multipart/related; boundary={mpwriter.boundary}"} + await self._async_upload(headers, backup_metadata, open_stream) + _LOGGER.debug( + "Uploaded backup: %s to: '%s'", + backup.backup_id, + backup_metadata["name"], + ) + except (ClientError, TimeoutError) as err: + _LOGGER.error("Upload backup error: %s", err) + raise BackupAgentError("Failed to upload backup") from err + + async with self._update_backups_json_lock: + backups_json_file_id, backups_json = await self._async_get_backups_json( + headers + ) + backups_json.append(backup.as_dict()) + await self._async_create_or_update_backups_json( + headers, backups_json_file_id, backups_json ) - try: - resp = await async_get_clientsession(self._hass).post( - DRIVE_API_UPLOAD_FILES, - params={"fields": ""}, - data=mpwriter, - headers=headers, - timeout=ClientTimeout(total=_UPLOAD_TIMEOUT), - ) - resp.raise_for_status() - except (ClientError, TimeoutError) as err: - raise BackupAgentError("Failed to upload backup") from err async def async_list_backups(self, **kwargs: Any) -> list[AgentBackup]: """List backups.""" headers = await self._async_headers() - try: - resp = await async_get_clientsession(self._hass).get( - DRIVE_API_FILES, - params={ - "q": f"'{self._folder_id}' in parents and trashed=false", - "fields": "files(properties)", - }, - headers=headers, - ) - resp.raise_for_status() - res = await resp.json() - except ClientError as err: - raise BackupAgentError("Failed to list backups") from err - backups = [] - for file in res["files"]: - if "properties" not in file or "backup_id" not in file["properties"]: - continue - backup = _convert_properties_to_agent_backup(file["properties"]) - backups.append(backup) - return backups + _, backups_json = await self._async_get_backups_json(headers) + return [AgentBackup.from_dict(backup) for backup in backups_json] async def async_get_backup( self, @@ -152,9 +135,11 @@ async def async_get_backup( **kwargs: Any, ) -> AgentBackup | None: """Return a backup.""" - headers = await self._async_headers() - _, backup = await self._async_get_file_id_and_properties(backup_id, headers) - return backup + backups = await self.async_list_backups() + for backup in backups: + if backup.backup_id == backup_id: + return backup + return None async def async_download_backup( self, @@ -166,20 +151,18 @@ async def async_download_backup( :param backup_id: The ID of the backup that was returned in async_list_backups. :return: An async iterator that yields bytes. """ + _LOGGER.debug("Downloading backup_id: %s", backup_id) headers = await self._async_headers() - file_id, _ = await self._async_get_file_id_and_properties(backup_id, headers) - if file_id is None: - raise BackupAgentError("Backup not found") - try: - resp = await async_get_clientsession(self._hass).get( - f"{DRIVE_API_FILES}/{file_id}", - params={"alt": "media"}, - headers=headers, - ) - resp.raise_for_status() - except ClientError as err: - raise BackupAgentError("Failed to download backup") from err - return ChunkAsyncStreamIterator(resp.content) + file_id = await self._async_get_backup_file_id(headers, backup_id) + if file_id: + try: + stream = await self._async_download(headers, file_id) + except ClientError as err: + _LOGGER.error("Download error: %s", err) + raise BackupAgentError("Failed to download backup") from err + return ChunkAsyncStreamIterator(stream) + _LOGGER.error("Download backup_id: %s not found", backup_id) + raise BackupAgentError("Backup not found") async def async_delete_backup( self, @@ -190,18 +173,28 @@ async def async_delete_backup( :param backup_id: The ID of the backup that was returned in async_list_backups. """ + _LOGGER.debug("Deleting backup_id: %s", backup_id) headers = await self._async_headers() - file_id, _ = await self._async_get_file_id_and_properties(backup_id, headers) - if file_id is None: - return - try: - resp = await async_get_clientsession(self._hass).delete( - f"{DRIVE_API_FILES}/{file_id}", - headers=headers, + file_id = await self._async_get_backup_file_id(headers, backup_id) + if file_id: + try: + resp = await async_get_clientsession(self._hass).delete( + f"{DRIVE_API_FILES}/{file_id}", + headers=headers, + ) + resp.raise_for_status() + _LOGGER.debug("Deleted backup_id: %s", backup_id) + except ClientError as err: + _LOGGER.error("Delete backup error: %s", err) + raise BackupAgentError("Failed to delete backup") from err + async with self._update_backups_json_lock: + backups_json_file_id, backups_json = await self._async_get_backups_json( + headers + ) + backups_json = [x for x in backups_json if x["backup_id"] != backup_id] + await self._async_create_or_update_backups_json( + headers, backups_json_file_id, backups_json ) - resp.raise_for_status() - except ClientError as err: - raise BackupAgentError("Failed to delete backup") from err async def _async_headers(self) -> dict[str, str]: try: @@ -210,28 +203,175 @@ async def _async_headers(self) -> dict[str, str]: raise BackupAgentError("Failed to refresh token") from err return create_headers(access_token) - async def _async_get_file_id_and_properties( - self, backup_id: str, headers: dict[str, str] - ) -> tuple[str | None, AgentBackup | None]: + async def _async_get_backups_json( + self, headers: dict[str, str] + ) -> tuple[str | None, list[dict[str, Any]]]: query = " and ".join( [ f"'{self._folder_id}' in parents", - f"properties has {{ key='backup_id' and value='{backup_id}' }}", + "trashed=false", + "properties has { key='ha' and value='backups.json' }", ] ) try: - resp = await async_get_clientsession(self._hass).get( - DRIVE_API_FILES, - params={ - "q": query, - "fields": "files(id,properties)", + res = await self._async_query(headers, query, "files(id)") + except ClientError as err: + _LOGGER.error("_async_get_backups_json error: %s", err) + raise BackupAgentError("Failed to get backups.json") from err + backups_json_file_id = None + files = res["files"] + for file in files: + backups_json_file_id = str(file["id"]) + if len(files) > 1: + _LOGGER.warning( + "Found multiple backups.json in %s/%s. Using %s", + DRIVE_FOLDER_URL_PREFIX, + self._folder_id, + backups_json_file_id, + ) + backups_json = [] + if backups_json_file_id: + all_bytes = bytearray() + try: + stream = await self._async_download(headers, backups_json_file_id) + except ClientError as err: + _LOGGER.error("_async_get_backups_json error: %s", err) + raise BackupAgentError("Failed to download backups.json") from err + async for chunk in stream: + all_bytes.extend(chunk) + backups_json = json.loads(all_bytes) + return backups_json_file_id, backups_json + + async def _async_create_or_update_backups_json( + self, + headers: dict[str, str], + backups_json_file_id: str | None, + backups_json: list[dict[str, Any]], + ) -> None: + def _create_open_stream_backup_json() -> Callable[[], Awaitable[bytes]]: + async def _open_stream_backup_json() -> bytes: + return json.dumps(backups_json, indent=2).encode("utf-8") + + return _open_stream_backup_json + + if backups_json_file_id: + try: + _LOGGER.debug("Updating backups.json") + await self._async_upload_existing( + headers, + backups_json_file_id, + _create_open_stream_backup_json(), + ) + _LOGGER.debug("Updated backups.json") + except ClientError as err: + _LOGGER.error("Update backups.json error: %s", err) + raise BackupAgentError("Failed to update backups.json") from err + else: + backups_json_metadata = { + "name": "backups.json", + "parents": [self._folder_id], + "properties": { + "ha": "backups.json", }, + } + try: + _LOGGER.debug("Creating backups.json") + await self._async_upload( + headers, + backups_json_metadata, + _create_open_stream_backup_json(), + ) + _LOGGER.debug("Created backups.json") + except ClientError as err: + _LOGGER.error("Create backups.json error: %s", err) + raise BackupAgentError("Failed to create backups.json") from err + + async def _async_upload( + self, + headers: dict[str, str], + file_metadata: dict[str, Any], + open_stream: Callable[ + [], Coroutine[Any, Any, AsyncIterator[bytes]] | Awaitable[bytes] + ], + ) -> None: + with MultipartWriter() as mpwriter: + mpwriter.append_json(file_metadata) + mpwriter.append(await open_stream()) + headers.update( + {"Content-Type": f"multipart/related; boundary={mpwriter.boundary}"} + ) + resp = await async_get_clientsession(self._hass).post( + DRIVE_API_UPLOAD_FILES, + params={"fields": ""}, + data=mpwriter, headers=headers, + timeout=ClientTimeout(total=_UPLOAD_TIMEOUT), ) resp.raise_for_status() - res = await resp.json() + + async def _async_upload_existing( + self, + headers: dict[str, str], + file_id: str, + open_stream: Callable[ + [], Coroutine[Any, Any, AsyncIterator[bytes]] | Awaitable[bytes] + ], + ) -> None: + resp = await async_get_clientsession(self._hass).patch( + f"{DRIVE_API_UPLOAD_FILES}/{file_id}", + params={"fields": ""}, + data=await open_stream(), + headers=headers, + ) + resp.raise_for_status() + + async def _async_download( + self, headers: dict[str, str], file_id: str + ) -> StreamReader: + resp = await async_get_clientsession(self._hass).get( + f"{DRIVE_API_FILES}/{file_id}", + params={"alt": "media"}, + headers=headers, + ) + resp.raise_for_status() + return resp.content + + async def _async_get_backup_file_id( + self, + headers: dict[str, str], + backup_id: str, + ) -> str | None: + query = " and ".join( + [ + f"'{self._folder_id}' in parents", + f"properties has {{ key='backup_id' and value='{backup_id}' }}", + ] + ) + try: + res = await self._async_query(headers, query, "files(id)") except ClientError as err: + _LOGGER.error("_async_get_backup_file_id error: %s", err) raise BackupAgentError("Failed to get backup") from err for file in res["files"]: - return file["id"], _convert_properties_to_agent_backup(file["properties"]) - return None, None + return str(file["id"]) + return None + + async def _async_query( + self, + headers: dict[str, str], + query: str, + fields: str, + ) -> dict[str, Any]: + _LOGGER.debug("_async_query: query: %s fields: %s", query, fields) + resp = await async_get_clientsession(self._hass).get( + DRIVE_API_FILES, + params={ + "q": query, + "fields": fields, + }, + headers=headers, + ) + resp.raise_for_status() + res: dict[str, Any] = await resp.json() + _LOGGER.debug("_async_query result: %s", res) + return res diff --git a/homeassistant/components/google_drive/config_flow.py b/homeassistant/components/google_drive/config_flow.py index 6a794dc3ae3eda..e3ae868a9ab6f7 100644 --- a/homeassistant/components/google_drive/config_flow.py +++ b/homeassistant/components/google_drive/config_flow.py @@ -77,9 +77,9 @@ async def async_oauth_create_entry(self, data: dict[str, Any]) -> ConfigFlowResu json={ "name": "Home Assistant", "mimeType": "application/vnd.google-apps.folder", - # Adding a property to be able to identify this folder - # if needed in the future. - "properties": {"ha": "root"}, + "properties": { + "ha": "root", + }, }, headers=headers, ) From 22e07cfb8b4fa9f5c723a557311966f71ce8e3fe Mon Sep 17 00:00:00 2001 From: tronikos Date: Wed, 8 Jan 2025 23:38:13 +0000 Subject: [PATCH 16/23] Address comments --- .../components/google_drive/strings.json | 2 +- tests/components/google_drive/conftest.py | 25 ++++++++++++++++++ .../google_drive/test_config_flow.py | 26 ++++--------------- tests/components/google_drive/test_init.py | 19 +++----------- 4 files changed, 35 insertions(+), 37 deletions(-) create mode 100644 tests/components/google_drive/conftest.py diff --git a/homeassistant/components/google_drive/strings.json b/homeassistant/components/google_drive/strings.json index f6d14219d6c864..608015fc72f79e 100644 --- a/homeassistant/components/google_drive/strings.json +++ b/homeassistant/components/google_drive/strings.json @@ -37,7 +37,7 @@ "message": "Google Drive folder not found. Remove the Google Drive integration and re-add it for a new 'Home Assistant' folder to be created in Google Drive." }, "config_entry_error_folder_4xx": { - "message": "Unexpected error. Open an issue." + "message": "Unexpected error while setting up Google Drive." } } } diff --git a/tests/components/google_drive/conftest.py b/tests/components/google_drive/conftest.py new file mode 100644 index 00000000000000..3f9eefd8825e0b --- /dev/null +++ b/tests/components/google_drive/conftest.py @@ -0,0 +1,25 @@ +"""PyTest fixtures and test helpers.""" + +import pytest + +from homeassistant.components.application_credentials import ( + ClientCredential, + async_import_client_credential, +) +from homeassistant.components.google_drive.const import DOMAIN +from homeassistant.core import HomeAssistant +from homeassistant.setup import async_setup_component + +CLIENT_ID = "1234" +CLIENT_SECRET = "5678" + + +@pytest.fixture +async def setup_credentials(hass: HomeAssistant) -> None: + """Fixture to setup credentials.""" + assert await async_setup_component(hass, "application_credentials", {}) + await async_import_client_credential( + hass, + DOMAIN, + ClientCredential(CLIENT_ID, CLIENT_SECRET), + ) diff --git a/tests/components/google_drive/test_config_flow.py b/tests/components/google_drive/test_config_flow.py index 6f51517deddce3..96bce8ff6307e3 100644 --- a/tests/components/google_drive/test_config_flow.py +++ b/tests/components/google_drive/test_config_flow.py @@ -6,39 +6,23 @@ import pytest from homeassistant import config_entries -from homeassistant.components.application_credentials import ( - ClientCredential, - async_import_client_credential, -) from homeassistant.components.google_drive.const import DOMAIN from homeassistant.core import HomeAssistant from homeassistant.data_entry_flow import FlowResultType from homeassistant.helpers import config_entry_oauth2_flow -from homeassistant.setup import async_setup_component + +from .conftest import CLIENT_ID from tests.common import MockConfigEntry from tests.test_util.aiohttp import AiohttpClientMocker from tests.typing import ClientSessionGenerator -CLIENT_ID = "1234" -CLIENT_SECRET = "5678" GOOGLE_AUTH_URI = "https://accounts.google.com/o/oauth2/v2/auth" GOOGLE_TOKEN_URI = "https://oauth2.googleapis.com/token" FOLDER_ID = "google-folder-id" TITLE = "Google Drive" -@pytest.fixture -async def setup_credentials(hass: HomeAssistant) -> None: - """Fixture to setup credentials.""" - assert await async_setup_component(hass, "application_credentials", {}) - await async_import_client_credential( - hass, - DOMAIN, - ClientCredential(CLIENT_ID, CLIENT_SECRET), - ) - - @pytest.mark.usefixtures("current_request_with_host") async def test_full_flow( hass: HomeAssistant, @@ -48,7 +32,7 @@ async def test_full_flow( ) -> None: """Check full flow.""" result = await hass.config_entries.flow.async_init( - "google_drive", context={"source": config_entries.SOURCE_USER} + DOMAIN, context={"source": config_entries.SOURCE_USER} ) state = config_entry_oauth2_flow._encode_jwt( hass, @@ -123,7 +107,7 @@ async def test_create_folder_error( ) -> None: """Test case where creating the folder fails.""" result = await hass.config_entries.flow.async_init( - "google_drive", context={"source": config_entries.SOURCE_USER} + DOMAIN, context={"source": config_entries.SOURCE_USER} ) state = config_entry_oauth2_flow._encode_jwt( hass, @@ -262,7 +246,7 @@ async def test_already_configured( config_entry.add_to_hass(hass) result = await hass.config_entries.flow.async_init( - "google_drive", context={"source": config_entries.SOURCE_USER} + DOMAIN, context={"source": config_entries.SOURCE_USER} ) assert result.get("type") is FlowResultType.ABORT assert result.get("reason") == "single_instance_allowed" diff --git a/tests/components/google_drive/test_init.py b/tests/components/google_drive/test_init.py index 714808a899c0bc..7a2c5d5d76a599 100644 --- a/tests/components/google_drive/test_init.py +++ b/tests/components/google_drive/test_init.py @@ -7,14 +7,9 @@ import pytest -from homeassistant.components.application_credentials import ( - ClientCredential, - async_import_client_credential, -) from homeassistant.components.google_drive import DOMAIN from homeassistant.config_entries import ConfigEntryState from homeassistant.core import HomeAssistant -from homeassistant.setup import async_setup_component from tests.common import MockConfigEntry from tests.test_util.aiohttp import AiohttpClientMocker @@ -50,21 +45,15 @@ def mock_config_entry(expires_at: int) -> MockConfigEntry: @pytest.fixture(name="setup_integration") async def mock_setup_integration( - hass: HomeAssistant, config_entry: MockConfigEntry + hass: HomeAssistant, + config_entry: MockConfigEntry, + setup_credentials, ) -> Callable[[], Coroutine[Any, Any, None]]: """Fixture for setting up the component.""" config_entry.add_to_hass(hass) - assert await async_setup_component(hass, "application_credentials", {}) - await async_import_client_credential( - hass, - DOMAIN, - ClientCredential("client-id", "client-secret"), - DOMAIN, - ) - async def func() -> None: - assert await async_setup_component(hass, DOMAIN, {}) + await hass.config_entries.async_setup(config_entry.entry_id) await hass.async_block_till_done() return func From 40856fc1817a4873224aea51e9938ea596386b40 Mon Sep 17 00:00:00 2001 From: tronikos Date: Fri, 10 Jan 2025 01:32:13 +0000 Subject: [PATCH 17/23] autouse=True on setup_credentials --- tests/components/google_drive/conftest.py | 2 +- tests/components/google_drive/test_config_flow.py | 4 ---- tests/components/google_drive/test_init.py | 1 - 3 files changed, 1 insertion(+), 6 deletions(-) diff --git a/tests/components/google_drive/conftest.py b/tests/components/google_drive/conftest.py index 3f9eefd8825e0b..2b03d955631035 100644 --- a/tests/components/google_drive/conftest.py +++ b/tests/components/google_drive/conftest.py @@ -14,7 +14,7 @@ CLIENT_SECRET = "5678" -@pytest.fixture +@pytest.fixture(autouse=True) async def setup_credentials(hass: HomeAssistant) -> None: """Fixture to setup credentials.""" assert await async_setup_component(hass, "application_credentials", {}) diff --git a/tests/components/google_drive/test_config_flow.py b/tests/components/google_drive/test_config_flow.py index 96bce8ff6307e3..ba2b8137dd14d1 100644 --- a/tests/components/google_drive/test_config_flow.py +++ b/tests/components/google_drive/test_config_flow.py @@ -28,7 +28,6 @@ async def test_full_flow( hass: HomeAssistant, hass_client_no_auth: ClientSessionGenerator, aioclient_mock: AiohttpClientMocker, - setup_credentials, ) -> None: """Check full flow.""" result = await hass.config_entries.flow.async_init( @@ -103,7 +102,6 @@ async def test_create_folder_error( hass: HomeAssistant, hass_client_no_auth: ClientSessionGenerator, aioclient_mock: AiohttpClientMocker, - setup_credentials, ) -> None: """Test case where creating the folder fails.""" result = await hass.config_entries.flow.async_init( @@ -155,7 +153,6 @@ async def test_reauth( hass: HomeAssistant, hass_client_no_auth: ClientSessionGenerator, aioclient_mock: AiohttpClientMocker, - setup_credentials, ) -> None: """Test the reauthentication case updates the existing config entry.""" @@ -231,7 +228,6 @@ async def test_already_configured( hass: HomeAssistant, hass_client_no_auth: ClientSessionGenerator, aioclient_mock: AiohttpClientMocker, - setup_credentials, ) -> None: """Test case for single_instance_allowed.""" config_entry = MockConfigEntry( diff --git a/tests/components/google_drive/test_init.py b/tests/components/google_drive/test_init.py index 7a2c5d5d76a599..df2be036df3eec 100644 --- a/tests/components/google_drive/test_init.py +++ b/tests/components/google_drive/test_init.py @@ -47,7 +47,6 @@ def mock_config_entry(expires_at: int) -> MockConfigEntry: async def mock_setup_integration( hass: HomeAssistant, config_entry: MockConfigEntry, - setup_credentials, ) -> Callable[[], Coroutine[Any, Any, None]]: """Fixture for setting up the component.""" config_entry.add_to_hass(hass) From 0bdf6a9b2acaeeae6f23de114c848a8cfdebadab Mon Sep 17 00:00:00 2001 From: tronikos Date: Fri, 10 Jan 2025 01:51:30 +0000 Subject: [PATCH 18/23] Store metadata in description and remove backups.json --- .../components/google_drive/backup.py | 147 +++--------------- 1 file changed, 21 insertions(+), 126 deletions(-) diff --git a/homeassistant/components/google_drive/backup.py b/homeassistant/components/google_drive/backup.py index eff067c14abeb0..2a8ba5cc98b51f 100644 --- a/homeassistant/components/google_drive/backup.py +++ b/homeassistant/components/google_drive/backup.py @@ -2,7 +2,6 @@ from __future__ import annotations -from asyncio import Lock from collections.abc import AsyncIterator, Awaitable, Callable, Coroutine import json import logging @@ -20,12 +19,7 @@ from . import DATA_BACKUP_AGENT_LISTENERS, GoogleDriveConfigEntry from .api import create_headers -from .const import ( - DOMAIN, - DRIVE_API_FILES, - DRIVE_API_UPLOAD_FILES, - DRIVE_FOLDER_URL_PREFIX, -) +from .const import DOMAIN, DRIVE_API_FILES, DRIVE_API_UPLOAD_FILES _LOGGER = logging.getLogger(__name__) _UPLOAD_TIMEOUT = 12 * 3600 @@ -75,7 +69,6 @@ def __init__( self._hass = hass self._folder_id = config_entry.unique_id self._auth = config_entry.runtime_data - self._update_backups_json_lock = Lock() async def async_upload_backup( self, @@ -92,6 +85,7 @@ async def async_upload_backup( headers = await self._async_headers() backup_metadata = { "name": f"{backup.name} {backup.date}.tar", + "description": json.dumps(backup.as_dict()), "parents": [self._folder_id], "properties": { "ha": "backup", @@ -114,20 +108,28 @@ async def async_upload_backup( _LOGGER.error("Upload backup error: %s", err) raise BackupAgentError("Failed to upload backup") from err - async with self._update_backups_json_lock: - backups_json_file_id, backups_json = await self._async_get_backups_json( - headers - ) - backups_json.append(backup.as_dict()) - await self._async_create_or_update_backups_json( - headers, backups_json_file_id, backups_json - ) - async def async_list_backups(self, **kwargs: Any) -> list[AgentBackup]: """List backups.""" headers = await self._async_headers() - _, backups_json = await self._async_get_backups_json(headers) - return [AgentBackup.from_dict(backup) for backup in backups_json] + query = " and ".join( + [ + f"'{self._folder_id}' in parents", + "trashed=false", + ] + ) + try: + res = await self._async_query( + headers, query, "files(description,properties)" + ) + except ClientError as err: + raise BackupAgentError("Failed to list backups") from err + backups = [] + for file in res["files"]: + if "properties" not in file or "backup_id" not in file["properties"]: + continue + backup = AgentBackup.from_dict(json.loads(file["description"])) + backups.append(backup) + return backups async def async_get_backup( self, @@ -187,14 +189,6 @@ async def async_delete_backup( except ClientError as err: _LOGGER.error("Delete backup error: %s", err) raise BackupAgentError("Failed to delete backup") from err - async with self._update_backups_json_lock: - backups_json_file_id, backups_json = await self._async_get_backups_json( - headers - ) - backups_json = [x for x in backups_json if x["backup_id"] != backup_id] - await self._async_create_or_update_backups_json( - headers, backups_json_file_id, backups_json - ) async def _async_headers(self) -> dict[str, str]: try: @@ -203,89 +197,6 @@ async def _async_headers(self) -> dict[str, str]: raise BackupAgentError("Failed to refresh token") from err return create_headers(access_token) - async def _async_get_backups_json( - self, headers: dict[str, str] - ) -> tuple[str | None, list[dict[str, Any]]]: - query = " and ".join( - [ - f"'{self._folder_id}' in parents", - "trashed=false", - "properties has { key='ha' and value='backups.json' }", - ] - ) - try: - res = await self._async_query(headers, query, "files(id)") - except ClientError as err: - _LOGGER.error("_async_get_backups_json error: %s", err) - raise BackupAgentError("Failed to get backups.json") from err - backups_json_file_id = None - files = res["files"] - for file in files: - backups_json_file_id = str(file["id"]) - if len(files) > 1: - _LOGGER.warning( - "Found multiple backups.json in %s/%s. Using %s", - DRIVE_FOLDER_URL_PREFIX, - self._folder_id, - backups_json_file_id, - ) - backups_json = [] - if backups_json_file_id: - all_bytes = bytearray() - try: - stream = await self._async_download(headers, backups_json_file_id) - except ClientError as err: - _LOGGER.error("_async_get_backups_json error: %s", err) - raise BackupAgentError("Failed to download backups.json") from err - async for chunk in stream: - all_bytes.extend(chunk) - backups_json = json.loads(all_bytes) - return backups_json_file_id, backups_json - - async def _async_create_or_update_backups_json( - self, - headers: dict[str, str], - backups_json_file_id: str | None, - backups_json: list[dict[str, Any]], - ) -> None: - def _create_open_stream_backup_json() -> Callable[[], Awaitable[bytes]]: - async def _open_stream_backup_json() -> bytes: - return json.dumps(backups_json, indent=2).encode("utf-8") - - return _open_stream_backup_json - - if backups_json_file_id: - try: - _LOGGER.debug("Updating backups.json") - await self._async_upload_existing( - headers, - backups_json_file_id, - _create_open_stream_backup_json(), - ) - _LOGGER.debug("Updated backups.json") - except ClientError as err: - _LOGGER.error("Update backups.json error: %s", err) - raise BackupAgentError("Failed to update backups.json") from err - else: - backups_json_metadata = { - "name": "backups.json", - "parents": [self._folder_id], - "properties": { - "ha": "backups.json", - }, - } - try: - _LOGGER.debug("Creating backups.json") - await self._async_upload( - headers, - backups_json_metadata, - _create_open_stream_backup_json(), - ) - _LOGGER.debug("Created backups.json") - except ClientError as err: - _LOGGER.error("Create backups.json error: %s", err) - raise BackupAgentError("Failed to create backups.json") from err - async def _async_upload( self, headers: dict[str, str], @@ -309,22 +220,6 @@ async def _async_upload( ) resp.raise_for_status() - async def _async_upload_existing( - self, - headers: dict[str, str], - file_id: str, - open_stream: Callable[ - [], Coroutine[Any, Any, AsyncIterator[bytes]] | Awaitable[bytes] - ], - ) -> None: - resp = await async_get_clientsession(self._hass).patch( - f"{DRIVE_API_UPLOAD_FILES}/{file_id}", - params={"fields": ""}, - data=await open_stream(), - headers=headers, - ) - resp.raise_for_status() - async def _async_download( self, headers: dict[str, str], file_id: str ) -> StreamReader: From fc54c610828f4103a824b72f5a2dedcb5891fbbb Mon Sep 17 00:00:00 2001 From: tronikos Date: Fri, 10 Jan 2025 11:39:39 +0000 Subject: [PATCH 19/23] improvements --- .../components/google_drive/__init__.py | 38 ++-- homeassistant/components/google_drive/api.py | 210 +++++++++++++++++- .../components/google_drive/backup.py | 161 ++------------ .../components/google_drive/config_flow.py | 63 +++--- .../components/google_drive/const.py | 8 - .../components/google_drive/manifest.json | 3 +- .../google_drive/quality_scale.yaml | 4 +- .../components/google_drive/strings.json | 8 +- .../google_drive/test_config_flow.py | 177 +++++++++++++-- tests/components/google_drive/test_init.py | 14 +- 10 files changed, 444 insertions(+), 242 deletions(-) diff --git a/homeassistant/components/google_drive/__init__.py b/homeassistant/components/google_drive/__init__.py index 729bb851b9c022..33e799084af04b 100644 --- a/homeassistant/components/google_drive/__init__.py +++ b/homeassistant/components/google_drive/__init__.py @@ -9,6 +9,7 @@ from homeassistant.config_entries import ConfigEntry from homeassistant.core import HomeAssistant from homeassistant.exceptions import ConfigEntryError, ConfigEntryNotReady +from homeassistant.helpers import instance_id from homeassistant.helpers.aiohttp_client import async_get_clientsession from homeassistant.helpers.config_entry_oauth2_flow import ( OAuth2Session, @@ -16,40 +17,43 @@ ) from homeassistant.util.hass_dict import HassKey -from .api import AsyncConfigEntryAuth, create_headers -from .const import DOMAIN, DRIVE_API_FILES +from .api import AsyncConfigEntryAuth, DriveClient +from .const import DOMAIN DATA_BACKUP_AGENT_LISTENERS: HassKey[list[Callable[[], None]]] = HassKey( f"{DOMAIN}.backup_agent_listeners" ) -type GoogleDriveConfigEntry = ConfigEntry[AsyncConfigEntryAuth] + +type GoogleDriveConfigEntry = ConfigEntry[DriveClient] async def async_setup_entry(hass: HomeAssistant, entry: GoogleDriveConfigEntry) -> bool: """Set up Google Drive from a config entry.""" implementation = await async_get_config_entry_implementation(hass, entry) - session = OAuth2Session(hass, entry, implementation) - auth = AsyncConfigEntryAuth(hass, session) - access_token = await auth.check_and_refresh_token() + auth = AsyncConfigEntryAuth(hass, OAuth2Session(hass, entry, implementation)) + + # Test we can refresh the token and raise ConfigEntryAuthFailed or ConfigEntryNotReady if not + await auth.check_and_refresh_token() + + client = DriveClient( + session=async_get_clientsession(hass), + ha_instance_id=await instance_id.async_get(hass), + access_token=None, + auth=auth, + ) + entry.runtime_data = client + + # Test we can access Google Drive and raise if not try: - resp = await async_get_clientsession(hass).get( - f"{DRIVE_API_FILES}/{entry.unique_id}", - params={"fields": ""}, - headers=create_headers(access_token), - ) - resp.raise_for_status() + await client.async_create_ha_root_folder_if_not_exists() except ClientError as err: if isinstance(err, ClientResponseError) and 400 <= err.status < 500: - if err.status == 404: - raise ConfigEntryError( - translation_key="config_entry_error_folder_not_found" - ) from err raise ConfigEntryError( translation_key="config_entry_error_folder_4xx" ) from err raise ConfigEntryNotReady from err - entry.runtime_data = auth + return True diff --git a/homeassistant/components/google_drive/api.py b/homeassistant/components/google_drive/api.py index 5328e3c2793300..2b91b615319db1 100644 --- a/homeassistant/components/google_drive/api.py +++ b/homeassistant/components/google_drive/api.py @@ -2,9 +2,16 @@ from __future__ import annotations +from collections.abc import AsyncIterator, Awaitable, Callable, Coroutine +import json +import logging +from typing import Any + +from aiohttp import ClientSession, ClientTimeout, MultipartWriter, StreamReader from aiohttp.client_exceptions import ClientError, ClientResponseError from google.auth.exceptions import RefreshError +from homeassistant.components.backup import AgentBackup from homeassistant.config_entries import ConfigEntryState from homeassistant.const import CONF_ACCESS_TOKEN from homeassistant.core import HomeAssistant @@ -15,12 +22,12 @@ ) from homeassistant.helpers import config_entry_oauth2_flow +DRIVE_API_ABOUT = "https://www.googleapis.com/drive/v3/about" +DRIVE_API_FILES = "https://www.googleapis.com/drive/v3/files" +DRIVE_API_UPLOAD_FILES = "https://www.googleapis.com/upload/drive/v3/files" +_UPLOAD_TIMEOUT = 12 * 3600 -def create_headers(access_token: str) -> dict[str, str]: - """Create headers with the provided access token.""" - return { - "Authorization": f"Bearer {access_token}", - } +_LOGGER = logging.getLogger(__name__) class AsyncConfigEntryAuth: @@ -64,3 +71,196 @@ async def check_and_refresh_token(self) -> str: ) raise HomeAssistantError(ex) from ex return self.access_token + + +class DriveClient: + """Google Drive client.""" + + def __init__( + self, + session: ClientSession, + ha_instance_id: str, + access_token: str | None, + auth: AsyncConfigEntryAuth | None, + ) -> None: + """Initialize Google Drive client.""" + self._session = session + self._ha_instance_id = ha_instance_id + self._access_token = access_token + self._auth = auth + assert self._access_token or self._auth + + async def _async_get_headers(self) -> dict[str, str]: + if self._access_token: + access_token = self._access_token + else: + assert self._auth + access_token = await self._auth.check_and_refresh_token() + return {"Authorization": f"Bearer {access_token}"} + + async def async_get_email_address(self) -> str: + """Get email address of the current user.""" + headers = await self._async_get_headers() + resp = await self._session.get( + DRIVE_API_ABOUT, + params={"fields": "user(emailAddress)"}, + headers=headers, + ) + resp.raise_for_status() + res = await resp.json() + return str(res["user"]["emailAddress"]) + + async def async_create_ha_root_folder_if_not_exists(self) -> tuple[str, str]: + """Create Home Assistant folder if it doesn't exist.""" + fields = "id,name" + query = " and ".join( + [ + "properties has { key='ha' and value='root' }", + f"properties has {{ key='instance_id' and value='{self._ha_instance_id}' }}", + ] + ) + res = await self.async_query(query, f"files({fields})") + for file in res["files"]: + _LOGGER.debug("Found existing folder: %s", file) + return str(file["id"]), str(file["name"]) + + headers = await self._async_get_headers() + file_metadata = { + "name": "Home Assistant", + "mimeType": "application/vnd.google-apps.folder", + "properties": { + "ha": "root", + "instance_id": self._ha_instance_id, + }, + } + _LOGGER.debug("Creating new folder with metadata: %s", file_metadata) + resp = await self._session.post( + DRIVE_API_FILES, + params={"fields": fields}, + json=file_metadata, + headers=headers, + ) + resp.raise_for_status() + res = await resp.json() + _LOGGER.debug("Created folder: %s", res) + return str(res["id"]), str(res["name"]) + + async def async_upload_backup( + self, + open_stream: Callable[[], Coroutine[Any, Any, AsyncIterator[bytes]]], + backup: AgentBackup, + ) -> None: + """Upload a backup.""" + folder_id, _ = await self.async_create_ha_root_folder_if_not_exists() + backup_metadata = { + "name": f"{backup.name} {backup.date}.tar", + "description": json.dumps(backup.as_dict()), + "parents": [folder_id], + "properties": { + "ha": "backup", + "instance_id": self._ha_instance_id, + "backup_id": backup.backup_id, + }, + } + _LOGGER.debug( + "Uploading backup: %s with Google Drive metadata: %s", + backup.backup_id, + backup_metadata, + ) + await self.async_upload(backup_metadata, open_stream) + _LOGGER.debug( + "Uploaded backup: %s to: '%s'", + backup.backup_id, + backup_metadata["name"], + ) + + async def async_list_backups(self) -> list[AgentBackup]: + """List backups.""" + query = " and ".join( + [ + "properties has { key='ha' and value='backup' }", + f"properties has {{ key='instance_id' and value='{self._ha_instance_id}' }}", + "trashed=false", + ] + ) + res = await self.async_query(query, "files(description)") + backups = [] + for file in res["files"]: + backup = AgentBackup.from_dict(json.loads(file["description"])) + backups.append(backup) + return backups + + async def async_get_backup_file_id(self, backup_id: str) -> str | None: + """Get file_id of backup if it exists.""" + query = " and ".join( + [ + "properties has { key='ha' and value='backup' }", + f"properties has {{ key='instance_id' and value='{self._ha_instance_id}' }}", + f"properties has {{ key='backup_id' and value='{backup_id}' }}", + ] + ) + res = await self.async_query(query, "files(id)") + for file in res["files"]: + return str(file["id"]) + return None + + async def async_delete(self, file_id: str) -> None: + """Delete file.""" + headers = await self._async_get_headers() + resp = await self._session.delete( + f"{DRIVE_API_FILES}/{file_id}", headers=headers + ) + resp.raise_for_status() + + async def async_download(self, file_id: str) -> StreamReader: + """Download a file.""" + headers = await self._async_get_headers() + resp = await self._session.get( + f"{DRIVE_API_FILES}/{file_id}", + params={"alt": "media"}, + headers=headers, + ) + resp.raise_for_status() + return resp.content + + async def async_upload( + self, + file_metadata: dict[str, Any], + open_stream: Callable[ + [], Coroutine[Any, Any, AsyncIterator[bytes]] | Awaitable[bytes] + ], + ) -> None: + """Upload a file.""" + headers = await self._async_get_headers() + with MultipartWriter() as mpwriter: + mpwriter.append_json(file_metadata) + mpwriter.append(await open_stream()) + headers.update( + {"Content-Type": f"multipart/related; boundary={mpwriter.boundary}"} + ) + resp = await self._session.post( + DRIVE_API_UPLOAD_FILES, + params={"fields": ""}, + data=mpwriter, + headers=headers, + timeout=ClientTimeout(total=_UPLOAD_TIMEOUT), + ) + resp.raise_for_status() + + async def async_query( + self, + query: str, + fields: str, + ) -> dict[str, Any]: + """Query for files.""" + headers = await self._async_get_headers() + _LOGGER.debug("async_query: query: '%s' fields: '%s'", query, fields) + resp = await self._session.get( + DRIVE_API_FILES, + params={"q": query, "fields": fields}, + headers=headers, + ) + resp.raise_for_status() + res: dict[str, Any] = await resp.json() + _LOGGER.debug("async_query result: %s", res) + return res diff --git a/homeassistant/components/google_drive/backup.py b/homeassistant/components/google_drive/backup.py index 2a8ba5cc98b51f..3e0ce303fffe4b 100644 --- a/homeassistant/components/google_drive/backup.py +++ b/homeassistant/components/google_drive/backup.py @@ -2,27 +2,20 @@ from __future__ import annotations -from collections.abc import AsyncIterator, Awaitable, Callable, Coroutine -import json +from collections.abc import AsyncIterator, Callable, Coroutine import logging from typing import Any -from aiohttp import ClientError, ClientTimeout, MultipartWriter, StreamReader +from aiohttp import ClientError from homeassistant.components.backup import AgentBackup, BackupAgent, BackupAgentError from homeassistant.core import HomeAssistant, callback -from homeassistant.exceptions import HomeAssistantError -from homeassistant.helpers.aiohttp_client import ( - ChunkAsyncStreamIterator, - async_get_clientsession, -) +from homeassistant.helpers.aiohttp_client import ChunkAsyncStreamIterator from . import DATA_BACKUP_AGENT_LISTENERS, GoogleDriveConfigEntry -from .api import create_headers -from .const import DOMAIN, DRIVE_API_FILES, DRIVE_API_UPLOAD_FILES +from .const import DOMAIN _LOGGER = logging.getLogger(__name__) -_UPLOAD_TIMEOUT = 12 * 3600 async def async_get_backup_agents( @@ -31,7 +24,7 @@ async def async_get_backup_agents( ) -> list[BackupAgent]: """Return a list of backup agents.""" entries = hass.config_entries.async_loaded_entries(DOMAIN) - return [GoogleDriveBackupAgent(hass, entry) for entry in entries] + return [GoogleDriveBackupAgent(entry) for entry in entries] @callback @@ -60,15 +53,11 @@ class GoogleDriveBackupAgent(BackupAgent): domain = DOMAIN - def __init__( - self, hass: HomeAssistant, config_entry: GoogleDriveConfigEntry - ) -> None: + def __init__(self, config_entry: GoogleDriveConfigEntry) -> None: """Initialize the cloud backup sync agent.""" super().__init__() self.name = config_entry.title - self._hass = hass - self._folder_id = config_entry.unique_id - self._auth = config_entry.runtime_data + self._client = config_entry.runtime_data async def async_upload_backup( self, @@ -82,54 +71,18 @@ async def async_upload_backup( :param open_stream: A function returning an async iterator that yields bytes. :param backup: Metadata about the backup that should be uploaded. """ - headers = await self._async_headers() - backup_metadata = { - "name": f"{backup.name} {backup.date}.tar", - "description": json.dumps(backup.as_dict()), - "parents": [self._folder_id], - "properties": { - "ha": "backup", - "backup_id": backup.backup_id, - }, - } try: - _LOGGER.debug( - "Uploading backup: %s with Google Drive metadata: %s", - backup, - backup_metadata, - ) - await self._async_upload(headers, backup_metadata, open_stream) - _LOGGER.debug( - "Uploaded backup: %s to: '%s'", - backup.backup_id, - backup_metadata["name"], - ) + await self._client.async_upload_backup(open_stream, backup) except (ClientError, TimeoutError) as err: _LOGGER.error("Upload backup error: %s", err) raise BackupAgentError("Failed to upload backup") from err async def async_list_backups(self, **kwargs: Any) -> list[AgentBackup]: """List backups.""" - headers = await self._async_headers() - query = " and ".join( - [ - f"'{self._folder_id}' in parents", - "trashed=false", - ] - ) try: - res = await self._async_query( - headers, query, "files(description,properties)" - ) + return await self._client.async_list_backups() except ClientError as err: raise BackupAgentError("Failed to list backups") from err - backups = [] - for file in res["files"]: - if "properties" not in file or "backup_id" not in file["properties"]: - continue - backup = AgentBackup.from_dict(json.loads(file["description"])) - backups.append(backup) - return backups async def async_get_backup( self, @@ -154,11 +107,11 @@ async def async_download_backup( :return: An async iterator that yields bytes. """ _LOGGER.debug("Downloading backup_id: %s", backup_id) - headers = await self._async_headers() - file_id = await self._async_get_backup_file_id(headers, backup_id) + file_id = await self._client.async_get_backup_file_id(backup_id) if file_id: + _LOGGER.debug("Downloading file_id: %s", file_id) try: - stream = await self._async_download(headers, file_id) + stream = await self._client.async_download(file_id) except ClientError as err: _LOGGER.error("Download error: %s", err) raise BackupAgentError("Failed to download backup") from err @@ -176,97 +129,11 @@ async def async_delete_backup( :param backup_id: The ID of the backup that was returned in async_list_backups. """ _LOGGER.debug("Deleting backup_id: %s", backup_id) - headers = await self._async_headers() - file_id = await self._async_get_backup_file_id(headers, backup_id) + file_id = await self._client.async_get_backup_file_id(backup_id) if file_id: try: - resp = await async_get_clientsession(self._hass).delete( - f"{DRIVE_API_FILES}/{file_id}", - headers=headers, - ) - resp.raise_for_status() + await self._client.async_delete(file_id) _LOGGER.debug("Deleted backup_id: %s", backup_id) except ClientError as err: _LOGGER.error("Delete backup error: %s", err) raise BackupAgentError("Failed to delete backup") from err - - async def _async_headers(self) -> dict[str, str]: - try: - access_token = await self._auth.check_and_refresh_token() - except HomeAssistantError as err: - raise BackupAgentError("Failed to refresh token") from err - return create_headers(access_token) - - async def _async_upload( - self, - headers: dict[str, str], - file_metadata: dict[str, Any], - open_stream: Callable[ - [], Coroutine[Any, Any, AsyncIterator[bytes]] | Awaitable[bytes] - ], - ) -> None: - with MultipartWriter() as mpwriter: - mpwriter.append_json(file_metadata) - mpwriter.append(await open_stream()) - headers.update( - {"Content-Type": f"multipart/related; boundary={mpwriter.boundary}"} - ) - resp = await async_get_clientsession(self._hass).post( - DRIVE_API_UPLOAD_FILES, - params={"fields": ""}, - data=mpwriter, - headers=headers, - timeout=ClientTimeout(total=_UPLOAD_TIMEOUT), - ) - resp.raise_for_status() - - async def _async_download( - self, headers: dict[str, str], file_id: str - ) -> StreamReader: - resp = await async_get_clientsession(self._hass).get( - f"{DRIVE_API_FILES}/{file_id}", - params={"alt": "media"}, - headers=headers, - ) - resp.raise_for_status() - return resp.content - - async def _async_get_backup_file_id( - self, - headers: dict[str, str], - backup_id: str, - ) -> str | None: - query = " and ".join( - [ - f"'{self._folder_id}' in parents", - f"properties has {{ key='backup_id' and value='{backup_id}' }}", - ] - ) - try: - res = await self._async_query(headers, query, "files(id)") - except ClientError as err: - _LOGGER.error("_async_get_backup_file_id error: %s", err) - raise BackupAgentError("Failed to get backup") from err - for file in res["files"]: - return str(file["id"]) - return None - - async def _async_query( - self, - headers: dict[str, str], - query: str, - fields: str, - ) -> dict[str, Any]: - _LOGGER.debug("_async_query: query: %s fields: %s", query, fields) - resp = await async_get_clientsession(self._hass).get( - DRIVE_API_FILES, - params={ - "q": query, - "fields": fields, - }, - headers=headers, - ) - resp.raise_for_status() - res: dict[str, Any] = await resp.json() - _LOGGER.debug("_async_query result: %s", res) - return res diff --git a/homeassistant/components/google_drive/config_flow.py b/homeassistant/components/google_drive/config_flow.py index e3ae868a9ab6f7..6785f6398b55a6 100644 --- a/homeassistant/components/google_drive/config_flow.py +++ b/homeassistant/components/google_drive/config_flow.py @@ -10,17 +10,17 @@ from homeassistant.config_entries import SOURCE_REAUTH, ConfigFlowResult from homeassistant.const import CONF_ACCESS_TOKEN, CONF_TOKEN -from homeassistant.helpers import config_entry_oauth2_flow +from homeassistant.helpers import config_entry_oauth2_flow, instance_id from homeassistant.helpers.aiohttp_client import async_get_clientsession -from .api import create_headers -from .const import ( - DEFAULT_NAME, - DOMAIN, - DRIVE_API_FILES, - DRIVE_FOLDER_URL_PREFIX, - OAUTH2_SCOPES, -) +from .api import DriveClient +from .const import DOMAIN + +DEFAULT_NAME = "Google Drive" +DRIVE_FOLDER_URL_PREFIX = "https://drive.google.com/drive/folders/" +OAUTH2_SCOPES = [ + "https://www.googleapis.com/auth/drive.file", +] class OAuth2FlowHandler( @@ -61,38 +61,43 @@ async def async_step_reauth_confirm( async def async_oauth_create_entry(self, data: dict[str, Any]) -> ConfigFlowResult: """Create an entry for the flow, or update existing entry.""" + client = DriveClient( + session=async_get_clientsession(self.hass), + ha_instance_id=await instance_id.async_get(self.hass), + access_token=data[CONF_TOKEN][CONF_ACCESS_TOKEN], + auth=None, + ) + + try: + email_address = await client.async_get_email_address() + except ClientError as err: + self.logger.error("Error getting email address: %s", err) + return self.async_abort(reason="cannot_connect") - session = async_get_clientsession(self.hass) - headers = create_headers(data[CONF_TOKEN][CONF_ACCESS_TOKEN]) + await self.async_set_unique_id(email_address) if self.source == SOURCE_REAUTH: + self._abort_if_unique_id_mismatch() return self.async_update_reload_and_abort( self._get_reauth_entry(), data=data ) + self._abort_if_unique_id_configured() + try: - resp = await session.post( - DRIVE_API_FILES, - params={"fields": "id"}, - json={ - "name": "Home Assistant", - "mimeType": "application/vnd.google-apps.folder", - "properties": { - "ha": "root", - }, - }, - headers=headers, - ) - resp.raise_for_status() - res = await resp.json() + ( + folder_id, + folder_name, + ) = await client.async_create_ha_root_folder_if_not_exists() except ClientError as err: self.logger.error("Error creating folder: %s", str(err)) return self.async_abort(reason="create_folder_failure") - folder_id = res["id"] - await self.async_set_unique_id(folder_id) - self._abort_if_unique_id_configured() + return self.async_create_entry( title=DEFAULT_NAME, data=data, - description_placeholders={"url": f"{DRIVE_FOLDER_URL_PREFIX}{folder_id}"}, + description_placeholders={ + "folder_name": folder_name, + "url": f"{DRIVE_FOLDER_URL_PREFIX}{folder_id}", + }, ) diff --git a/homeassistant/components/google_drive/const.py b/homeassistant/components/google_drive/const.py index 39eab5b4acf9cc..3f0b3e9d6100a0 100644 --- a/homeassistant/components/google_drive/const.py +++ b/homeassistant/components/google_drive/const.py @@ -3,11 +3,3 @@ from __future__ import annotations DOMAIN = "google_drive" - -DEFAULT_NAME = "Google Drive" -DRIVE_API_FILES = "https://www.googleapis.com/drive/v3/files" -DRIVE_API_UPLOAD_FILES = "https://www.googleapis.com/upload/drive/v3/files" -DRIVE_FOLDER_URL_PREFIX = "https://drive.google.com/drive/folders/" -OAUTH2_SCOPES = [ - "https://www.googleapis.com/auth/drive.file", -] diff --git a/homeassistant/components/google_drive/manifest.json b/homeassistant/components/google_drive/manifest.json index 32231b16616a88..4bff9101480131 100644 --- a/homeassistant/components/google_drive/manifest.json +++ b/homeassistant/components/google_drive/manifest.json @@ -8,6 +8,5 @@ "documentation": "https://www.home-assistant.io/integrations/google_drive", "integration_type": "service", "iot_class": "cloud_polling", - "quality_scale": "bronze", - "single_config_entry": true + "quality_scale": "bronze" } diff --git a/homeassistant/components/google_drive/quality_scale.yaml b/homeassistant/components/google_drive/quality_scale.yaml index 7f7b7bd76e1462..3bbd6081bc4aa7 100644 --- a/homeassistant/components/google_drive/quality_scale.yaml +++ b/homeassistant/components/google_drive/quality_scale.yaml @@ -100,7 +100,9 @@ rules: reconfiguration-flow: status: exempt comment: No configuration options. - repair-issues: todo + repair-issues: + status: exempt + comment: No repairs. stale-devices: status: exempt comment: No devices. diff --git a/homeassistant/components/google_drive/strings.json b/homeassistant/components/google_drive/strings.json index 608015fc72f79e..f993f4ffd70c73 100644 --- a/homeassistant/components/google_drive/strings.json +++ b/homeassistant/components/google_drive/strings.json @@ -23,19 +23,17 @@ "create_folder_failure": "Error while creating Google Drive folder, see error log for details", "oauth_timeout": "[%key:common::config_flow::abort::oauth2_timeout%]", "oauth_unauthorized": "[%key:common::config_flow::abort::oauth2_unauthorized%]", - "oauth_failed": "[%key:common::config_flow::abort::oauth2_failed%]" + "oauth_failed": "[%key:common::config_flow::abort::oauth2_failed%]", + "unique_id_mismatch": "The new email address does not match the previous one" }, "create_entry": { - "default": "Successfully authenticated and created 'Home Assistant' folder at: {url}\nFeel free to rename it in Google Drive as you wish." + "default": "Using [{folder_name}]({url}) folder. Feel free to rename it in Google Drive as you wish." } }, "application_credentials": { "description": "Follow the [instructions]({more_info_url}) for [OAuth consent screen]({oauth_consent_url}) to give Home Assistant access to your Google Drive. You also need to create Application Credentials linked to your account:\n1. Go to [Credentials]({oauth_creds_url}) and select **Create Credentials**.\n1. From the drop-down list select **OAuth client ID**.\n1. Select **Web application** for the Application Type." }, "exceptions": { - "config_entry_error_folder_not_found": { - "message": "Google Drive folder not found. Remove the Google Drive integration and re-add it for a new 'Home Assistant' folder to be created in Google Drive." - }, "config_entry_error_folder_4xx": { "message": "Unexpected error while setting up Google Drive." } diff --git a/tests/components/google_drive/test_config_flow.py b/tests/components/google_drive/test_config_flow.py index ba2b8137dd14d1..697f88c69a376c 100644 --- a/tests/components/google_drive/test_config_flow.py +++ b/tests/components/google_drive/test_config_flow.py @@ -9,7 +9,7 @@ from homeassistant.components.google_drive.const import DOMAIN from homeassistant.core import HomeAssistant from homeassistant.data_entry_flow import FlowResultType -from homeassistant.helpers import config_entry_oauth2_flow +from homeassistant.helpers import config_entry_oauth2_flow, instance_id from .conftest import CLIENT_ID @@ -19,7 +19,9 @@ GOOGLE_AUTH_URI = "https://accounts.google.com/o/oauth2/v2/auth" GOOGLE_TOKEN_URI = "https://oauth2.googleapis.com/token" +USER_EMAIL = "user@domain.com" FOLDER_ID = "google-folder-id" +FOLDER_NAME = "folder name" TITLE = "Google Drive" @@ -53,10 +55,18 @@ async def test_full_flow( assert resp.status == 200 assert resp.headers["content-type"] == "text/html; charset=utf-8" - # Prepare API response when creating the folder + # Prepare API responses + aioclient_mock.get( + "https://www.googleapis.com/drive/v3/about", + json={"user": {"emailAddress": USER_EMAIL}}, + ) + aioclient_mock.get( + "https://www.googleapis.com/drive/v3/files", + json={"files": []}, + ) aioclient_mock.post( - "https://www.googleapis.com/drive/v3/files?fields=id", - json={"id": FOLDER_ID}, + "https://www.googleapis.com/drive/v3/files", + json={"id": FOLDER_ID, "name": FOLDER_NAME}, ) aioclient_mock.post( @@ -76,20 +86,23 @@ async def test_full_flow( assert len(hass.config_entries.async_entries(DOMAIN)) == 1 assert len(mock_setup.mock_calls) == 1 - assert len(aioclient_mock.mock_calls) == 2 - assert aioclient_mock.mock_calls[1][2] == { + assert len(aioclient_mock.mock_calls) == 4 + assert aioclient_mock.mock_calls[3][2] == { "name": "Home Assistant", "mimeType": "application/vnd.google-apps.folder", - "properties": {"ha": "root"}, + "properties": { + "ha": "root", + "instance_id": await instance_id.async_get(hass), + }, } - assert aioclient_mock.mock_calls[1][3] == { + assert aioclient_mock.mock_calls[3][3] == { "Authorization": "Bearer mock-access-token" } assert result.get("type") is FlowResultType.CREATE_ENTRY assert result.get("title") == TITLE assert "result" in result - assert result.get("result").unique_id == FOLDER_ID + assert result.get("result").unique_id == USER_EMAIL assert "token" in result.get("result").data assert result.get("result").data["token"].get("access_token") == "mock-access-token" assert ( @@ -127,9 +140,17 @@ async def test_create_folder_error( assert resp.status == 200 assert resp.headers["content-type"] == "text/html; charset=utf-8" - # Prepare fake exception creating the folder + # Prepare API responses + aioclient_mock.get( + "https://www.googleapis.com/drive/v3/about", + json={"user": {"emailAddress": USER_EMAIL}}, + ) + aioclient_mock.get( + "https://www.googleapis.com/drive/v3/files", + json={"files": []}, + ) aioclient_mock.post( - "https://www.googleapis.com/drive/v3/files?fields=id", + "https://www.googleapis.com/drive/v3/files", exc=ClientError, ) @@ -149,16 +170,82 @@ async def test_create_folder_error( @pytest.mark.usefixtures("current_request_with_host") +async def test_get_email_error( + hass: HomeAssistant, + hass_client_no_auth: ClientSessionGenerator, + aioclient_mock: AiohttpClientMocker, +) -> None: + """Test case where getting the email address fails.""" + result = await hass.config_entries.flow.async_init( + DOMAIN, context={"source": config_entries.SOURCE_USER} + ) + state = config_entry_oauth2_flow._encode_jwt( + hass, + { + "flow_id": result["flow_id"], + "redirect_uri": "https://example.com/auth/external/callback", + }, + ) + + assert result["url"] == ( + f"{GOOGLE_AUTH_URI}?response_type=code&client_id={CLIENT_ID}" + "&redirect_uri=https://example.com/auth/external/callback" + f"&state={state}&scope=https://www.googleapis.com/auth/drive.file" + "&access_type=offline&prompt=consent" + ) + + client = await hass_client_no_auth() + resp = await client.get(f"/auth/external/callback?code=abcd&state={state}") + assert resp.status == 200 + assert resp.headers["content-type"] == "text/html; charset=utf-8" + + # Prepare API responses + aioclient_mock.get( + "https://www.googleapis.com/drive/v3/about", + exc=ClientError, + ) + aioclient_mock.post( + GOOGLE_TOKEN_URI, + json={ + "refresh_token": "mock-refresh-token", + "access_token": "mock-access-token", + "type": "Bearer", + "expires_in": 60, + }, + ) + + result = await hass.config_entries.flow.async_configure(result["flow_id"]) + assert result.get("type") is FlowResultType.ABORT + assert result.get("reason") == "cannot_connect" + + +@pytest.mark.usefixtures("current_request_with_host") +@pytest.mark.parametrize( + ("new_email", "expected_abort_reason"), + [ + ( + USER_EMAIL, + "reauth_successful", + ), + ( + "other.user@domain.com", + "unique_id_mismatch", + ), + ], + ids=["reauth_successful", "unique_id_mismatch"], +) async def test_reauth( hass: HomeAssistant, hass_client_no_auth: ClientSessionGenerator, aioclient_mock: AiohttpClientMocker, + new_email: str, + expected_abort_reason: str, ) -> None: - """Test the reauthentication case updates the existing config entry.""" + """Test the reauthentication flow.""" config_entry = MockConfigEntry( domain=DOMAIN, - unique_id=FOLDER_ID, + unique_id=USER_EMAIL, data={ "token": { "access_token": "mock-access-token", @@ -194,6 +281,11 @@ async def test_reauth( assert resp.status == 200 assert resp.headers["content-type"] == "text/html; charset=utf-8" + # Prepare API responses + aioclient_mock.get( + "https://www.googleapis.com/drive/v3/about", + json={"user": {"emailAddress": new_email}}, + ) aioclient_mock.post( GOOGLE_TOKEN_URI, json={ @@ -211,16 +303,23 @@ async def test_reauth( await hass.async_block_till_done() assert len(hass.config_entries.async_entries(DOMAIN)) == 1 - assert len(mock_setup.mock_calls) == 1 + if expected_abort_reason == "reauth_successful": + assert len(mock_setup.mock_calls) == 1 + else: + assert len(mock_setup.mock_calls) == 0 assert result.get("type") is FlowResultType.ABORT - assert result.get("reason") == "reauth_successful" + assert result.get("reason") == expected_abort_reason - assert config_entry.unique_id == FOLDER_ID + assert config_entry.unique_id == USER_EMAIL assert "token" in config_entry.data # Verify access token is refreshed - assert config_entry.data["token"].get("access_token") == "updated-access-token" - assert config_entry.data["token"].get("refresh_token") == "mock-refresh-token" + if expected_abort_reason == "reauth_successful": + assert config_entry.data["token"].get("access_token") == "updated-access-token" + assert config_entry.data["token"].get("refresh_token") == "mock-refresh-token" + else: + assert config_entry.data["token"].get("access_token") == "mock-access-token" + assert config_entry.data["token"].get("refresh_token") is None @pytest.mark.usefixtures("current_request_with_host") @@ -229,10 +328,10 @@ async def test_already_configured( hass_client_no_auth: ClientSessionGenerator, aioclient_mock: AiohttpClientMocker, ) -> None: - """Test case for single_instance_allowed.""" + """Test already configured account.""" config_entry = MockConfigEntry( domain=DOMAIN, - unique_id=FOLDER_ID, + unique_id=USER_EMAIL, data={ "token": { "access_token": "mock-access-token", @@ -244,5 +343,41 @@ async def test_already_configured( result = await hass.config_entries.flow.async_init( DOMAIN, context={"source": config_entries.SOURCE_USER} ) + state = config_entry_oauth2_flow._encode_jwt( + hass, + { + "flow_id": result["flow_id"], + "redirect_uri": "https://example.com/auth/external/callback", + }, + ) + + assert result["url"] == ( + f"{GOOGLE_AUTH_URI}?response_type=code&client_id={CLIENT_ID}" + "&redirect_uri=https://example.com/auth/external/callback" + f"&state={state}&scope=https://www.googleapis.com/auth/drive.file" + "&access_type=offline&prompt=consent" + ) + + client = await hass_client_no_auth() + resp = await client.get(f"/auth/external/callback?code=abcd&state={state}") + assert resp.status == 200 + assert resp.headers["content-type"] == "text/html; charset=utf-8" + + # Prepare API responses + aioclient_mock.get( + "https://www.googleapis.com/drive/v3/about", + json={"user": {"emailAddress": USER_EMAIL}}, + ) + aioclient_mock.post( + GOOGLE_TOKEN_URI, + json={ + "refresh_token": "mock-refresh-token", + "access_token": "mock-access-token", + "type": "Bearer", + "expires_in": 60, + }, + ) + + result = await hass.config_entries.flow.async_configure(result["flow_id"]) assert result.get("type") is FlowResultType.ABORT - assert result.get("reason") == "single_instance_allowed" + assert result.get("reason") == "already_configured" diff --git a/tests/components/google_drive/test_init.py b/tests/components/google_drive/test_init.py index df2be036df3eec..09ee0f3a731338 100644 --- a/tests/components/google_drive/test_init.py +++ b/tests/components/google_drive/test_init.py @@ -14,7 +14,7 @@ from tests.common import MockConfigEntry from tests.test_util.aiohttp import AiohttpClientMocker -TEST_FOLDER_ID = "google-folder-it" +TEST_USER_EMAIL = "testuser@domain.com" type ComponentSetup = Callable[[], Awaitable[None]] @@ -30,7 +30,7 @@ def mock_config_entry(expires_at: int) -> MockConfigEntry: """Fixture for MockConfigEntry.""" return MockConfigEntry( domain=DOMAIN, - unique_id=TEST_FOLDER_ID, + unique_id=TEST_USER_EMAIL, data={ "auth_implementation": DOMAIN, "token": { @@ -66,8 +66,8 @@ async def test_setup_success( """Test successful setup and unload.""" # Setup looks up existing folder to make sure it still exists aioclient_mock.get( - f"https://www.googleapis.com/drive/v3/files/{TEST_FOLDER_ID}?fields=", - json={}, + "https://www.googleapis.com/drive/v3/files", + json={"files": [{"id": "HA folder ID", "name": "HA folder name"}]}, ) await setup_integration() @@ -112,7 +112,7 @@ async def test_setup_error( """Test setup error.""" # Simulate failure looking up existing folder aioclient_mock.get( - f"https://www.googleapis.com/drive/v3/files/{TEST_FOLDER_ID}?fields=", + "https://www.googleapis.com/drive/v3/files", status=status, ) @@ -132,8 +132,8 @@ async def test_expired_token_refresh_success( """Test expired token is refreshed.""" # Setup looks up existing folder to make sure it still exists aioclient_mock.get( - f"https://www.googleapis.com/drive/v3/files/{TEST_FOLDER_ID}?fields=", - json={}, + "https://www.googleapis.com/drive/v3/files", + json={"files": [{"id": "HA folder ID", "name": "HA folder name"}]}, ) aioclient_mock.post( From e08f6c9d73be0ec1ef94c49921eb634abd7d7231 Mon Sep 17 00:00:00 2001 From: tronikos Date: Fri, 10 Jan 2025 12:49:42 +0000 Subject: [PATCH 20/23] timeout downloads --- homeassistant/components/google_drive/api.py | 5 +++-- homeassistant/components/google_drive/backup.py | 2 +- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/homeassistant/components/google_drive/api.py b/homeassistant/components/google_drive/api.py index 2b91b615319db1..5fffd0e72e2830 100644 --- a/homeassistant/components/google_drive/api.py +++ b/homeassistant/components/google_drive/api.py @@ -25,7 +25,7 @@ DRIVE_API_ABOUT = "https://www.googleapis.com/drive/v3/about" DRIVE_API_FILES = "https://www.googleapis.com/drive/v3/files" DRIVE_API_UPLOAD_FILES = "https://www.googleapis.com/upload/drive/v3/files" -_UPLOAD_TIMEOUT = 12 * 3600 +_UPLOAD_AND_DOWNLOAD_TIMEOUT = 12 * 3600 _LOGGER = logging.getLogger(__name__) @@ -219,6 +219,7 @@ async def async_download(self, file_id: str) -> StreamReader: f"{DRIVE_API_FILES}/{file_id}", params={"alt": "media"}, headers=headers, + timeout=ClientTimeout(total=_UPLOAD_AND_DOWNLOAD_TIMEOUT), ) resp.raise_for_status() return resp.content @@ -243,7 +244,7 @@ async def async_upload( params={"fields": ""}, data=mpwriter, headers=headers, - timeout=ClientTimeout(total=_UPLOAD_TIMEOUT), + timeout=ClientTimeout(total=_UPLOAD_AND_DOWNLOAD_TIMEOUT), ) resp.raise_for_status() diff --git a/homeassistant/components/google_drive/backup.py b/homeassistant/components/google_drive/backup.py index 3e0ce303fffe4b..7b4f50ad62f613 100644 --- a/homeassistant/components/google_drive/backup.py +++ b/homeassistant/components/google_drive/backup.py @@ -112,7 +112,7 @@ async def async_download_backup( _LOGGER.debug("Downloading file_id: %s", file_id) try: stream = await self._client.async_download(file_id) - except ClientError as err: + except (ClientError, TimeoutError) as err: _LOGGER.error("Download error: %s", err) raise BackupAgentError("Failed to download backup") from err return ChunkAsyncStreamIterator(stream) From 4efee898468d11c32e3f1ef0a8e2193fedc4f667 Mon Sep 17 00:00:00 2001 From: tronikos Date: Sat, 11 Jan 2025 12:10:36 +0000 Subject: [PATCH 21/23] library --- .../components/google_drive/__init__.py | 27 +- homeassistant/components/google_drive/api.py | 163 ++----- .../components/google_drive/backup.py | 11 +- .../components/google_drive/config_flow.py | 28 +- .../components/google_drive/manifest.json | 4 +- .../google_drive/quality_scale.yaml | 2 +- .../components/google_drive/strings.json | 23 +- requirements_all.txt | 3 + requirements_test_all.txt | 3 + tests/components/google_drive/conftest.py | 24 + .../google_drive/snapshots/test_backup.ambr | 159 ++++++ .../snapshots/test_config_flow.ambr | 44 ++ tests/components/google_drive/test_backup.py | 454 ++++++++++++++++++ .../google_drive/test_config_flow.py | 84 ++-- tests/components/google_drive/test_init.py | 47 +- 15 files changed, 824 insertions(+), 252 deletions(-) create mode 100644 tests/components/google_drive/snapshots/test_backup.ambr create mode 100644 tests/components/google_drive/snapshots/test_config_flow.ambr create mode 100644 tests/components/google_drive/test_backup.py diff --git a/homeassistant/components/google_drive/__init__.py b/homeassistant/components/google_drive/__init__.py index 33e799084af04b..4777db31443ec7 100644 --- a/homeassistant/components/google_drive/__init__.py +++ b/homeassistant/components/google_drive/__init__.py @@ -4,11 +4,11 @@ from collections.abc import Callable -from aiohttp.client_exceptions import ClientError, ClientResponseError +from google_drive_api.exceptions import GoogleDriveApiError from homeassistant.config_entries import ConfigEntry from homeassistant.core import HomeAssistant -from homeassistant.exceptions import ConfigEntryError, ConfigEntryNotReady +from homeassistant.exceptions import ConfigEntryNotReady from homeassistant.helpers import instance_id from homeassistant.helpers.aiohttp_client import async_get_clientsession from homeassistant.helpers.config_entry_oauth2_flow import ( @@ -30,28 +30,23 @@ async def async_setup_entry(hass: HomeAssistant, entry: GoogleDriveConfigEntry) -> bool: """Set up Google Drive from a config entry.""" - implementation = await async_get_config_entry_implementation(hass, entry) - auth = AsyncConfigEntryAuth(hass, OAuth2Session(hass, entry, implementation)) + auth = AsyncConfigEntryAuth( + async_get_clientsession(hass), + OAuth2Session( + hass, entry, await async_get_config_entry_implementation(hass, entry) + ), + ) # Test we can refresh the token and raise ConfigEntryAuthFailed or ConfigEntryNotReady if not - await auth.check_and_refresh_token() + await auth.async_get_access_token() - client = DriveClient( - session=async_get_clientsession(hass), - ha_instance_id=await instance_id.async_get(hass), - access_token=None, - auth=auth, - ) + client = DriveClient(await instance_id.async_get(hass), auth) entry.runtime_data = client # Test we can access Google Drive and raise if not try: await client.async_create_ha_root_folder_if_not_exists() - except ClientError as err: - if isinstance(err, ClientResponseError) and 400 <= err.status < 500: - raise ConfigEntryError( - translation_key="config_entry_error_folder_4xx" - ) from err + except GoogleDriveApiError as err: raise ConfigEntryNotReady from err return True diff --git a/homeassistant/components/google_drive/api.py b/homeassistant/components/google_drive/api.py index 5fffd0e72e2830..347631dc395f2e 100644 --- a/homeassistant/components/google_drive/api.py +++ b/homeassistant/components/google_drive/api.py @@ -2,19 +2,19 @@ from __future__ import annotations -from collections.abc import AsyncIterator, Awaitable, Callable, Coroutine +from collections.abc import AsyncIterator, Callable, Coroutine import json import logging from typing import Any -from aiohttp import ClientSession, ClientTimeout, MultipartWriter, StreamReader +from aiohttp import ClientSession, ClientTimeout, StreamReader from aiohttp.client_exceptions import ClientError, ClientResponseError from google.auth.exceptions import RefreshError +from google_drive_api.api import AbstractAuth, GoogleDriveApi from homeassistant.components.backup import AgentBackup from homeassistant.config_entries import ConfigEntryState from homeassistant.const import CONF_ACCESS_TOKEN -from homeassistant.core import HomeAssistant from homeassistant.exceptions import ( ConfigEntryAuthFailed, ConfigEntryNotReady, @@ -22,38 +22,30 @@ ) from homeassistant.helpers import config_entry_oauth2_flow -DRIVE_API_ABOUT = "https://www.googleapis.com/drive/v3/about" -DRIVE_API_FILES = "https://www.googleapis.com/drive/v3/files" -DRIVE_API_UPLOAD_FILES = "https://www.googleapis.com/upload/drive/v3/files" _UPLOAD_AND_DOWNLOAD_TIMEOUT = 12 * 3600 _LOGGER = logging.getLogger(__name__) -class AsyncConfigEntryAuth: +class AsyncConfigEntryAuth(AbstractAuth): """Provide Google Drive authentication tied to an OAuth2 based config entry.""" def __init__( self, - hass: HomeAssistant, - oauth2_session: config_entry_oauth2_flow.OAuth2Session, + websession: ClientSession, + oauth_session: config_entry_oauth2_flow.OAuth2Session, ) -> None: - """Initialize Google Drive Auth.""" - self._hass = hass - self.oauth_session = oauth2_session + """Initialize AsyncConfigEntryAuth.""" + super().__init__(websession) + self._oauth_session = oauth_session - @property - def access_token(self) -> str: - """Return the access token.""" - return str(self.oauth_session.token[CONF_ACCESS_TOKEN]) - - async def check_and_refresh_token(self) -> str: - """Check the token.""" + async def async_get_access_token(self) -> str: + """Return a valid access token.""" try: - await self.oauth_session.async_ensure_token_valid() + await self._oauth_session.async_ensure_token_valid() except (RefreshError, ClientResponseError, ClientError) as ex: if ( - self.oauth_session.config_entry.state + self._oauth_session.config_entry.state is ConfigEntryState.SETUP_IN_PROGRESS ): if isinstance(ex, ClientResponseError) and 400 <= ex.status < 500: @@ -66,11 +58,28 @@ async def check_and_refresh_token(self) -> str: or hasattr(ex, "status") and ex.status == 400 ): - self.oauth_session.config_entry.async_start_reauth( - self.oauth_session.hass + self._oauth_session.config_entry.async_start_reauth( + self._oauth_session.hass ) raise HomeAssistantError(ex) from ex - return self.access_token + return str(self._oauth_session.token[CONF_ACCESS_TOKEN]) + + +class AsyncConfigFlowAuth(AbstractAuth): + """Provide authentication tied to a fixed token for the config flow.""" + + def __init__( + self, + websession: ClientSession, + token: str, + ) -> None: + """Initialize AsyncConfigFlowAuth.""" + super().__init__(websession) + self._token = token + + async def async_get_access_token(self) -> str: + """Return a valid access token.""" + return self._token class DriveClient: @@ -78,36 +87,16 @@ class DriveClient: def __init__( self, - session: ClientSession, ha_instance_id: str, - access_token: str | None, - auth: AsyncConfigEntryAuth | None, + auth: AbstractAuth, ) -> None: """Initialize Google Drive client.""" - self._session = session self._ha_instance_id = ha_instance_id - self._access_token = access_token - self._auth = auth - assert self._access_token or self._auth - - async def _async_get_headers(self) -> dict[str, str]: - if self._access_token: - access_token = self._access_token - else: - assert self._auth - access_token = await self._auth.check_and_refresh_token() - return {"Authorization": f"Bearer {access_token}"} + self._api = GoogleDriveApi(auth) async def async_get_email_address(self) -> str: """Get email address of the current user.""" - headers = await self._async_get_headers() - resp = await self._session.get( - DRIVE_API_ABOUT, - params={"fields": "user(emailAddress)"}, - headers=headers, - ) - resp.raise_for_status() - res = await resp.json() + res = await self._api.get_user(params={"fields": "user(emailAddress)"}) return str(res["user"]["emailAddress"]) async def async_create_ha_root_folder_if_not_exists(self) -> tuple[str, str]: @@ -117,14 +106,16 @@ async def async_create_ha_root_folder_if_not_exists(self) -> tuple[str, str]: [ "properties has { key='ha' and value='root' }", f"properties has {{ key='instance_id' and value='{self._ha_instance_id}' }}", + "trashed=false", ] ) - res = await self.async_query(query, f"files({fields})") + res = await self._api.list_files( + params={"q": query, "fields": f"files({fields})"} + ) for file in res["files"]: _LOGGER.debug("Found existing folder: %s", file) return str(file["id"]), str(file["name"]) - headers = await self._async_get_headers() file_metadata = { "name": "Home Assistant", "mimeType": "application/vnd.google-apps.folder", @@ -134,14 +125,7 @@ async def async_create_ha_root_folder_if_not_exists(self) -> tuple[str, str]: }, } _LOGGER.debug("Creating new folder with metadata: %s", file_metadata) - resp = await self._session.post( - DRIVE_API_FILES, - params={"fields": fields}, - json=file_metadata, - headers=headers, - ) - resp.raise_for_status() - res = await resp.json() + res = await self._api.create_file(params={"fields": fields}, json=file_metadata) _LOGGER.debug("Created folder: %s", res) return str(res["id"]), str(res["name"]) @@ -167,7 +151,7 @@ async def async_upload_backup( backup.backup_id, backup_metadata, ) - await self.async_upload(backup_metadata, open_stream) + await self._api.upload_file(backup_metadata, open_stream) _LOGGER.debug( "Uploaded backup: %s to: '%s'", backup.backup_id, @@ -183,7 +167,9 @@ async def async_list_backups(self) -> list[AgentBackup]: "trashed=false", ] ) - res = await self.async_query(query, "files(description)") + res = await self._api.list_files( + params={"q": query, "fields": "files(description)"} + ) backups = [] for file in res["files"]: backup = AgentBackup.from_dict(json.loads(file["description"])) @@ -199,69 +185,18 @@ async def async_get_backup_file_id(self, backup_id: str) -> str | None: f"properties has {{ key='backup_id' and value='{backup_id}' }}", ] ) - res = await self.async_query(query, "files(id)") + res = await self._api.list_files(params={"q": query, "fields": "files(id)"}) for file in res["files"]: return str(file["id"]) return None async def async_delete(self, file_id: str) -> None: """Delete file.""" - headers = await self._async_get_headers() - resp = await self._session.delete( - f"{DRIVE_API_FILES}/{file_id}", headers=headers - ) - resp.raise_for_status() + await self._api.delete_file(file_id) async def async_download(self, file_id: str) -> StreamReader: """Download a file.""" - headers = await self._async_get_headers() - resp = await self._session.get( - f"{DRIVE_API_FILES}/{file_id}", - params={"alt": "media"}, - headers=headers, - timeout=ClientTimeout(total=_UPLOAD_AND_DOWNLOAD_TIMEOUT), + resp = await self._api.get_file_content( + file_id, timeout=ClientTimeout(total=_UPLOAD_AND_DOWNLOAD_TIMEOUT) ) - resp.raise_for_status() return resp.content - - async def async_upload( - self, - file_metadata: dict[str, Any], - open_stream: Callable[ - [], Coroutine[Any, Any, AsyncIterator[bytes]] | Awaitable[bytes] - ], - ) -> None: - """Upload a file.""" - headers = await self._async_get_headers() - with MultipartWriter() as mpwriter: - mpwriter.append_json(file_metadata) - mpwriter.append(await open_stream()) - headers.update( - {"Content-Type": f"multipart/related; boundary={mpwriter.boundary}"} - ) - resp = await self._session.post( - DRIVE_API_UPLOAD_FILES, - params={"fields": ""}, - data=mpwriter, - headers=headers, - timeout=ClientTimeout(total=_UPLOAD_AND_DOWNLOAD_TIMEOUT), - ) - resp.raise_for_status() - - async def async_query( - self, - query: str, - fields: str, - ) -> dict[str, Any]: - """Query for files.""" - headers = await self._async_get_headers() - _LOGGER.debug("async_query: query: '%s' fields: '%s'", query, fields) - resp = await self._session.get( - DRIVE_API_FILES, - params={"q": query, "fields": fields}, - headers=headers, - ) - resp.raise_for_status() - res: dict[str, Any] = await resp.json() - _LOGGER.debug("async_query result: %s", res) - return res diff --git a/homeassistant/components/google_drive/backup.py b/homeassistant/components/google_drive/backup.py index 7b4f50ad62f613..fc998a7e641353 100644 --- a/homeassistant/components/google_drive/backup.py +++ b/homeassistant/components/google_drive/backup.py @@ -6,7 +6,7 @@ import logging from typing import Any -from aiohttp import ClientError +from google_drive_api.exceptions import GoogleDriveApiError from homeassistant.components.backup import AgentBackup, BackupAgent, BackupAgentError from homeassistant.core import HomeAssistant, callback @@ -73,7 +73,7 @@ async def async_upload_backup( """ try: await self._client.async_upload_backup(open_stream, backup) - except (ClientError, TimeoutError) as err: + except (GoogleDriveApiError, TimeoutError) as err: _LOGGER.error("Upload backup error: %s", err) raise BackupAgentError("Failed to upload backup") from err @@ -81,7 +81,7 @@ async def async_list_backups(self, **kwargs: Any) -> list[AgentBackup]: """List backups.""" try: return await self._client.async_list_backups() - except ClientError as err: + except GoogleDriveApiError as err: raise BackupAgentError("Failed to list backups") from err async def async_get_backup( @@ -112,7 +112,7 @@ async def async_download_backup( _LOGGER.debug("Downloading file_id: %s", file_id) try: stream = await self._client.async_download(file_id) - except (ClientError, TimeoutError) as err: + except (GoogleDriveApiError, TimeoutError) as err: _LOGGER.error("Download error: %s", err) raise BackupAgentError("Failed to download backup") from err return ChunkAsyncStreamIterator(stream) @@ -131,9 +131,10 @@ async def async_delete_backup( _LOGGER.debug("Deleting backup_id: %s", backup_id) file_id = await self._client.async_get_backup_file_id(backup_id) if file_id: + _LOGGER.debug("Deleting file_id: %s", file_id) try: await self._client.async_delete(file_id) _LOGGER.debug("Deleted backup_id: %s", backup_id) - except ClientError as err: + except GoogleDriveApiError as err: _LOGGER.error("Delete backup error: %s", err) raise BackupAgentError("Failed to delete backup") from err diff --git a/homeassistant/components/google_drive/config_flow.py b/homeassistant/components/google_drive/config_flow.py index 6785f6398b55a6..ea945d466fc0f1 100644 --- a/homeassistant/components/google_drive/config_flow.py +++ b/homeassistant/components/google_drive/config_flow.py @@ -6,14 +6,14 @@ import logging from typing import Any -from aiohttp.client_exceptions import ClientError +from google_drive_api.exceptions import GoogleDriveApiError from homeassistant.config_entries import SOURCE_REAUTH, ConfigFlowResult from homeassistant.const import CONF_ACCESS_TOKEN, CONF_TOKEN from homeassistant.helpers import config_entry_oauth2_flow, instance_id from homeassistant.helpers.aiohttp_client import async_get_clientsession -from .api import DriveClient +from .api import AsyncConfigFlowAuth, DriveClient from .const import DOMAIN DEFAULT_NAME = "Google Drive" @@ -62,22 +62,25 @@ async def async_step_reauth_confirm( async def async_oauth_create_entry(self, data: dict[str, Any]) -> ConfigFlowResult: """Create an entry for the flow, or update existing entry.""" client = DriveClient( - session=async_get_clientsession(self.hass), - ha_instance_id=await instance_id.async_get(self.hass), - access_token=data[CONF_TOKEN][CONF_ACCESS_TOKEN], - auth=None, + await instance_id.async_get(self.hass), + AsyncConfigFlowAuth( + async_get_clientsession(self.hass), data[CONF_TOKEN][CONF_ACCESS_TOKEN] + ), ) try: email_address = await client.async_get_email_address() - except ClientError as err: + except GoogleDriveApiError as err: self.logger.error("Error getting email address: %s", err) - return self.async_abort(reason="cannot_connect") + return self.async_abort( + reason="access_not_configured", + description_placeholders={"message": str(err)}, + ) await self.async_set_unique_id(email_address) if self.source == SOURCE_REAUTH: - self._abort_if_unique_id_mismatch() + self._abort_if_unique_id_mismatch(reason="wrong_account") return self.async_update_reload_and_abort( self._get_reauth_entry(), data=data ) @@ -89,9 +92,12 @@ async def async_oauth_create_entry(self, data: dict[str, Any]) -> ConfigFlowResu folder_id, folder_name, ) = await client.async_create_ha_root_folder_if_not_exists() - except ClientError as err: + except GoogleDriveApiError as err: self.logger.error("Error creating folder: %s", str(err)) - return self.async_abort(reason="create_folder_failure") + return self.async_abort( + reason="create_folder_failure", + description_placeholders={"message": str(err)}, + ) return self.async_create_entry( title=DEFAULT_NAME, diff --git a/homeassistant/components/google_drive/manifest.json b/homeassistant/components/google_drive/manifest.json index 4bff9101480131..a1abb9b260a4bf 100644 --- a/homeassistant/components/google_drive/manifest.json +++ b/homeassistant/components/google_drive/manifest.json @@ -8,5 +8,7 @@ "documentation": "https://www.home-assistant.io/integrations/google_drive", "integration_type": "service", "iot_class": "cloud_polling", - "quality_scale": "bronze" + "loggers": ["google_drive_api"], + "quality_scale": "platinum", + "requirements": ["python-google-drive-api==0.0.2"] } diff --git a/homeassistant/components/google_drive/quality_scale.yaml b/homeassistant/components/google_drive/quality_scale.yaml index 3bbd6081bc4aa7..70627a6a6d7d6a 100644 --- a/homeassistant/components/google_drive/quality_scale.yaml +++ b/homeassistant/components/google_drive/quality_scale.yaml @@ -49,7 +49,7 @@ rules: status: exempt comment: No actions and no entities. reauthentication-flow: done - test-coverage: todo + test-coverage: done # Gold devices: diff --git a/homeassistant/components/google_drive/strings.json b/homeassistant/components/google_drive/strings.json index f993f4ffd70c73..093e2674e4030b 100644 --- a/homeassistant/components/google_drive/strings.json +++ b/homeassistant/components/google_drive/strings.json @@ -3,28 +3,24 @@ "step": { "pick_implementation": { "title": "[%key:common::config_flow::title::oauth2_pick_implementation%]" - }, - "auth": { "title": "Link Google Account" }, - "reauth_confirm": { - "title": "[%key:common::config_flow::title::reauth%]", - "description": "The Google Drive integration needs to re-authenticate your account" } }, "abort": { "already_configured": "[%key:common::config_flow::abort::already_configured_account%]", "already_in_progress": "[%key:common::config_flow::abort::already_in_progress%]", - "cannot_connect": "[%key:common::config_flow::error::cannot_connect%]", - "timeout_connect": "[%key:common::config_flow::error::timeout_connect%]", "oauth_error": "[%key:common::config_flow::abort::oauth2_error%]", "missing_configuration": "[%key:common::config_flow::abort::oauth2_missing_configuration%]", - "reauth_successful": "[%key:common::config_flow::abort::reauth_successful%]", - "invalid_access_token": "[%key:common::config_flow::error::invalid_access_token%]", + "authorize_url_timeout": "[%key:common::config_flow::abort::oauth2_authorize_url_timeout%]", + "no_url_available": "[%key:common::config_flow::abort::oauth2_no_url_available%]", + "user_rejected_authorize": "[%key:common::config_flow::abort::oauth2_user_rejected_authorize%]", "unknown": "[%key:common::config_flow::error::unknown%]", - "create_folder_failure": "Error while creating Google Drive folder, see error log for details", "oauth_timeout": "[%key:common::config_flow::abort::oauth2_timeout%]", "oauth_unauthorized": "[%key:common::config_flow::abort::oauth2_unauthorized%]", "oauth_failed": "[%key:common::config_flow::abort::oauth2_failed%]", - "unique_id_mismatch": "The new email address does not match the previous one" + "access_not_configured": "Unable to access the Google Drive API:\n\n{message}", + "create_folder_failure": "Error while creating Google Drive folder:\n\n{message}", + "wrong_account": "Wrong account: Please authenticate with the right account.", + "reauth_successful": "[%key:common::config_flow::abort::reauth_successful%]" }, "create_entry": { "default": "Using [{folder_name}]({url}) folder. Feel free to rename it in Google Drive as you wish." @@ -32,10 +28,5 @@ }, "application_credentials": { "description": "Follow the [instructions]({more_info_url}) for [OAuth consent screen]({oauth_consent_url}) to give Home Assistant access to your Google Drive. You also need to create Application Credentials linked to your account:\n1. Go to [Credentials]({oauth_creds_url}) and select **Create Credentials**.\n1. From the drop-down list select **OAuth client ID**.\n1. Select **Web application** for the Application Type." - }, - "exceptions": { - "config_entry_error_folder_4xx": { - "message": "Unexpected error while setting up Google Drive." - } } } diff --git a/requirements_all.txt b/requirements_all.txt index 6835acdd782831..5d85de639b930c 100644 --- a/requirements_all.txt +++ b/requirements_all.txt @@ -2374,6 +2374,9 @@ python-gc100==1.0.3a0 # homeassistant.components.gitlab_ci python-gitlab==1.6.0 +# homeassistant.components.google_drive +python-google-drive-api==0.0.2 + # homeassistant.components.analytics_insights python-homeassistant-analytics==0.8.1 diff --git a/requirements_test_all.txt b/requirements_test_all.txt index 5728eae5333a90..3abe65da957161 100644 --- a/requirements_test_all.txt +++ b/requirements_test_all.txt @@ -1916,6 +1916,9 @@ python-fullykiosk==0.0.14 # homeassistant.components.sms # python-gammu==3.2.4 +# homeassistant.components.google_drive +python-google-drive-api==0.0.2 + # homeassistant.components.analytics_insights python-homeassistant-analytics==0.8.1 diff --git a/tests/components/google_drive/conftest.py b/tests/components/google_drive/conftest.py index 2b03d955631035..494e4d88fa8aaa 100644 --- a/tests/components/google_drive/conftest.py +++ b/tests/components/google_drive/conftest.py @@ -1,5 +1,8 @@ """PyTest fixtures and test helpers.""" +from collections.abc import Generator +from unittest.mock import AsyncMock, MagicMock, patch + import pytest from homeassistant.components.application_credentials import ( @@ -12,6 +15,7 @@ CLIENT_ID = "1234" CLIENT_SECRET = "5678" +HA_UUID = "0a123c" @pytest.fixture(autouse=True) @@ -23,3 +27,23 @@ async def setup_credentials(hass: HomeAssistant) -> None: DOMAIN, ClientCredential(CLIENT_ID, CLIENT_SECRET), ) + + +@pytest.fixture +def mock_api() -> Generator[MagicMock]: + """Return a mocked GoogleDriveApi.""" + with patch( + "homeassistant.components.google_drive.api.GoogleDriveApi" + ) as mock_api_cl: + mock_api = mock_api_cl.return_value + yield mock_api + + +@pytest.fixture(autouse=True) +def mock_instance_id() -> Generator[AsyncMock]: + """Mock instance_id.""" + with patch( + "homeassistant.components.google_drive.config_flow.instance_id.async_get", + ) as mock_async_get: + mock_async_get.return_value = HA_UUID + yield mock_async_get diff --git a/tests/components/google_drive/snapshots/test_backup.ambr b/tests/components/google_drive/snapshots/test_backup.ambr new file mode 100644 index 00000000000000..b535096a3109ff --- /dev/null +++ b/tests/components/google_drive/snapshots/test_backup.ambr @@ -0,0 +1,159 @@ +# serializer version: 1 +# name: test_agents_delete + list([ + tuple( + 'list_files', + tuple( + ), + dict({ + 'params': dict({ + 'fields': 'files(id,name)', + 'q': "properties has { key='ha' and value='root' } and properties has { key='instance_id' and value='0a123c' } and trashed=false", + }), + }), + ), + tuple( + 'list_files', + tuple( + ), + dict({ + 'params': dict({ + 'fields': 'files(id)', + 'q': "properties has { key='ha' and value='backup' } and properties has { key='instance_id' and value='0a123c' } and properties has { key='backup_id' and value='test-backup' }", + }), + }), + ), + tuple( + 'delete_file', + tuple( + 'backup-file-id', + ), + dict({ + }), + ), + ]) +# --- +# name: test_agents_download + list([ + tuple( + 'list_files', + tuple( + ), + dict({ + 'params': dict({ + 'fields': 'files(id,name)', + 'q': "properties has { key='ha' and value='root' } and properties has { key='instance_id' and value='0a123c' } and trashed=false", + }), + }), + ), + tuple( + 'list_files', + tuple( + ), + dict({ + 'params': dict({ + 'fields': 'files(description)', + 'q': "properties has { key='ha' and value='backup' } and properties has { key='instance_id' and value='0a123c' } and trashed=false", + }), + }), + ), + tuple( + 'list_files', + tuple( + ), + dict({ + 'params': dict({ + 'fields': 'files(id)', + 'q': "properties has { key='ha' and value='backup' } and properties has { key='instance_id' and value='0a123c' } and properties has { key='backup_id' and value='test-backup' }", + }), + }), + ), + tuple( + 'get_file_content', + tuple( + 'backup-file-id', + ), + dict({ + 'timeout': dict({ + 'ceil_threshold': 5, + 'connect': None, + 'sock_connect': None, + 'sock_read': None, + 'total': 43200, + }), + }), + ), + ]) +# --- +# name: test_agents_list_backups + list([ + tuple( + 'list_files', + tuple( + ), + dict({ + 'params': dict({ + 'fields': 'files(id,name)', + 'q': "properties has { key='ha' and value='root' } and properties has { key='instance_id' and value='0a123c' } and trashed=false", + }), + }), + ), + tuple( + 'list_files', + tuple( + ), + dict({ + 'params': dict({ + 'fields': 'files(description)', + 'q': "properties has { key='ha' and value='backup' } and properties has { key='instance_id' and value='0a123c' } and trashed=false", + }), + }), + ), + ]) +# --- +# name: test_agents_upload + list([ + tuple( + 'list_files', + tuple( + ), + dict({ + 'params': dict({ + 'fields': 'files(id,name)', + 'q': "properties has { key='ha' and value='root' } and properties has { key='instance_id' and value='0a123c' } and trashed=false", + }), + }), + ), + tuple( + 'list_files', + tuple( + ), + dict({ + 'params': dict({ + 'fields': 'files(id,name)', + 'q': "properties has { key='ha' and value='root' } and properties has { key='instance_id' and value='0a123c' } and trashed=false", + }), + }), + ), + tuple( + 'upload_file', + tuple( + dict({ + 'description': '{"addons": [{"name": "Test", "slug": "test", "version": "1.0.0"}], "backup_id": "test-backup", "date": "2025-01-01T01:23:45.678Z", "database_included": true, "extra_metadata": {"with_automatic_settings": false}, "folders": [], "homeassistant_included": true, "homeassistant_version": "2024.12.0", "name": "Test", "protected": false, "size": 987}', + 'name': 'Test 2025-01-01T01:23:45.678Z.tar', + 'parents': list([ + 'HA folder ID', + ]), + 'properties': dict({ + 'backup_id': 'test-backup', + 'ha': 'backup', + 'instance_id': '0a123c', + }), + }), + "CoreBackupReaderWriter.async_receive_backup..open_backup() -> 'AsyncIterator[bytes]'", + ), + dict({ + }), + ), + ]) +# --- diff --git a/tests/components/google_drive/snapshots/test_config_flow.ambr b/tests/components/google_drive/snapshots/test_config_flow.ambr new file mode 100644 index 00000000000000..8615d289f09560 --- /dev/null +++ b/tests/components/google_drive/snapshots/test_config_flow.ambr @@ -0,0 +1,44 @@ +# serializer version: 1 +# name: test_full_flow + list([ + tuple( + 'get_user', + tuple( + ), + dict({ + 'params': dict({ + 'fields': 'user(emailAddress)', + }), + }), + ), + tuple( + 'list_files', + tuple( + ), + dict({ + 'params': dict({ + 'fields': 'files(id,name)', + 'q': "properties has { key='ha' and value='root' } and properties has { key='instance_id' and value='0a123c' } and trashed=false", + }), + }), + ), + tuple( + 'create_file', + tuple( + ), + dict({ + 'json': dict({ + 'mimeType': 'application/vnd.google-apps.folder', + 'name': 'Home Assistant', + 'properties': dict({ + 'ha': 'root', + 'instance_id': '0a123c', + }), + }), + 'params': dict({ + 'fields': 'id,name', + }), + }), + ), + ]) +# --- diff --git a/tests/components/google_drive/test_backup.py b/tests/components/google_drive/test_backup.py new file mode 100644 index 00000000000000..1f885f20575d89 --- /dev/null +++ b/tests/components/google_drive/test_backup.py @@ -0,0 +1,454 @@ +"""Test the Google Drive backup platform.""" + +from io import StringIO +import json +import time +from typing import Any +from unittest.mock import AsyncMock, MagicMock, Mock, patch + +from aiohttp import ClientResponse +from google_drive_api.exceptions import GoogleDriveApiError +import pytest +from syrupy.assertion import SnapshotAssertion + +from homeassistant.components.application_credentials import ( + ClientCredential, + async_import_client_credential, +) +from homeassistant.components.backup import ( + DOMAIN as BACKUP_DOMAIN, + AddonInfo, + AgentBackup, +) +from homeassistant.components.google_drive import DOMAIN +from homeassistant.core import HomeAssistant +from homeassistant.setup import async_setup_component + +from tests.common import MockConfigEntry +from tests.test_util.aiohttp import mock_stream +from tests.typing import ClientSessionGenerator, WebSocketGenerator + +FOLDER_ID = "google-folder-it" +TEST_USER_EMAIL = "testuser@domain.com" +CONFIG_ENTRY_TITLE = "Google Drive entry title" +TEST_AGENT_BACKUP = AgentBackup( + addons=[AddonInfo(name="Test", slug="test", version="1.0.0")], + backup_id="test-backup", + database_included=True, + date="2025-01-01T01:23:45.678Z", + extra_metadata={ + "with_automatic_settings": False, + }, + folders=[], + homeassistant_included=True, + homeassistant_version="2024.12.0", + name="Test", + protected=False, + size=987, +) + + +@pytest.fixture(name="config_entry") +def mock_config_entry() -> MockConfigEntry: + """Fixture for MockConfigEntry.""" + return MockConfigEntry( + domain=DOMAIN, + unique_id=TEST_USER_EMAIL, + title=CONFIG_ENTRY_TITLE, + data={ + "auth_implementation": DOMAIN, + "token": { + "access_token": "mock-access-token", + "refresh_token": "mock-refresh-token", + "expires_at": time.time() + 3600, + "scope": "https://www.googleapis.com/auth/drive.file", + }, + }, + ) + + +@pytest.fixture(autouse=True) +async def setup_integration( + hass: HomeAssistant, + config_entry: MockConfigEntry, + mock_api: MagicMock, +) -> None: + """Set up Google Drive integration.""" + config_entry.add_to_hass(hass) + assert await async_setup_component(hass, BACKUP_DOMAIN, {BACKUP_DOMAIN: {}}) + assert await async_setup_component(hass, "application_credentials", {}) + await async_import_client_credential( + hass, + DOMAIN, + ClientCredential("client-id", "client-secret"), + DOMAIN, + ) + mock_api.list_files = AsyncMock( + return_value={"files": [{"id": "HA folder ID", "name": "HA folder name"}]} + ) + assert await async_setup_component(hass, DOMAIN, {}) + await hass.async_block_till_done() + + +async def test_agents_info( + hass: HomeAssistant, + hass_ws_client: WebSocketGenerator, +) -> 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"}, + {"agent_id": f"google_drive.{CONFIG_ENTRY_TITLE}"}, + ], + } + + config_entry = hass.config_entries.async_entries(DOMAIN)[0] + await hass.config_entries.async_unload(config_entry.entry_id) + await hass.async_block_till_done() + + 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"}]} + + +async def test_agents_list_backups( + hass: HomeAssistant, + hass_ws_client: WebSocketGenerator, + mock_api: MagicMock, + snapshot: SnapshotAssertion, +) -> None: + """Test agent list backups.""" + mock_api.list_files = AsyncMock( + return_value={ + "files": [{"description": json.dumps(TEST_AGENT_BACKUP.as_dict())}] + } + ) + + 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"] == [ + TEST_AGENT_BACKUP.as_frontend_json() + | { + "agent_ids": [f"google_drive.{CONFIG_ENTRY_TITLE}"], + "failed_agent_ids": [], + "with_automatic_settings": None, + } + ] + assert [tuple(mock_call) for mock_call in mock_api.mock_calls] == snapshot + + +async def test_agents_list_backups_fail( + hass: HomeAssistant, + hass_ws_client: WebSocketGenerator, + mock_api: MagicMock, +) -> None: + """Test agent list backups fails.""" + mock_api.list_files = AsyncMock(side_effect=GoogleDriveApiError("some error")) + + 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": { + f"google_drive.{CONFIG_ENTRY_TITLE}": "Failed to list backups" + }, + "backups": [], + "last_attempted_automatic_backup": None, + "last_completed_automatic_backup": None, + } + + +@pytest.mark.parametrize( + ("backup_id", "expected_result"), + [ + ( + TEST_AGENT_BACKUP.backup_id, + TEST_AGENT_BACKUP.as_frontend_json() + | { + "agent_ids": [f"google_drive.{CONFIG_ENTRY_TITLE}"], + "failed_agent_ids": [], + "with_automatic_settings": None, + }, + ), + ( + "12345", + None, + ), + ], + ids=["found", "not_found"], +) +async def test_agents_get_backup( + hass: HomeAssistant, + hass_ws_client: WebSocketGenerator, + mock_api: MagicMock, + backup_id: str, + expected_result: dict[str, Any] | None, +) -> None: + """Test agent get backup.""" + mock_api.list_files = AsyncMock( + return_value={ + "files": [{"description": json.dumps(TEST_AGENT_BACKUP.as_dict())}] + } + ) + client = await hass_ws_client(hass) + await client.send_json_auto_id({"type": "backup/details", "backup_id": backup_id}) + response = await client.receive_json() + + assert response["success"] + assert response["result"]["agent_errors"] == {} + assert response["result"]["backup"] == expected_result + + +async def test_agents_download( + hass: HomeAssistant, + hass_client: ClientSessionGenerator, + mock_api: MagicMock, + snapshot: SnapshotAssertion, +) -> None: + """Test agent download backup.""" + mock_api.list_files = AsyncMock( + side_effect=[ + {"files": [{"description": json.dumps(TEST_AGENT_BACKUP.as_dict())}]}, + {"files": [{"id": "backup-file-id"}]}, + ] + ) + mock_response = AsyncMock(spec=ClientResponse) + mock_response.content = mock_stream(b"backup data") + mock_api.get_file_content = AsyncMock(return_value=mock_response) + + client = await hass_client() + resp = await client.get( + f"/api/backup/download/{TEST_AGENT_BACKUP.backup_id}?agent_id=google_drive.{CONFIG_ENTRY_TITLE}" + ) + assert resp.status == 200 + assert await resp.content.read() == b"backup data" + + assert [tuple(mock_call) for mock_call in mock_api.mock_calls] == snapshot + + +async def test_agents_download_fail( + hass: HomeAssistant, + hass_client: ClientSessionGenerator, + mock_api: MagicMock, +) -> None: + """Test agent download backup fails.""" + mock_api.list_files = AsyncMock( + side_effect=[ + {"files": [{"description": json.dumps(TEST_AGENT_BACKUP.as_dict())}]}, + {"files": [{"id": "backup-file-id"}]}, + ] + ) + mock_response = AsyncMock(spec=ClientResponse) + mock_response.content = mock_stream(b"backup data") + mock_api.get_file_content = AsyncMock(side_effect=GoogleDriveApiError("some error")) + + client = await hass_client() + resp = await client.get( + f"/api/backup/download/{TEST_AGENT_BACKUP.backup_id}?agent_id=google_drive.{CONFIG_ENTRY_TITLE}" + ) + assert resp.status == 500 + content = await resp.content.read() + assert "Failed to download backup" in content.decode() + + +async def test_agents_download_file_not_found( + hass: HomeAssistant, + hass_client: ClientSessionGenerator, + mock_api: MagicMock, +) -> None: + """Test agent download backup raises error if not found.""" + mock_api.list_files = AsyncMock( + side_effect=[ + {"files": [{"description": json.dumps(TEST_AGENT_BACKUP.as_dict())}]}, + {"files": []}, + ] + ) + + client = await hass_client() + resp = await client.get( + f"/api/backup/download/{TEST_AGENT_BACKUP.backup_id}?agent_id=google_drive.{CONFIG_ENTRY_TITLE}" + ) + assert resp.status == 500 + content = await resp.content.read() + assert "Backup not found" in content.decode() + + +async def test_agents_download_metadata_not_found( + hass: HomeAssistant, + hass_client: ClientSessionGenerator, + mock_api: MagicMock, +) -> None: + """Test agent download backup raises error if not found.""" + mock_api.list_files = AsyncMock( + return_value={ + "files": [{"description": json.dumps(TEST_AGENT_BACKUP.as_dict())}] + } + ) + + client = await hass_client() + backup_id = "1234" + assert backup_id != TEST_AGENT_BACKUP.backup_id + + resp = await client.get( + f"/api/backup/download/{backup_id}?agent_id=google_drive.{CONFIG_ENTRY_TITLE}" + ) + assert resp.status == 404 + assert await resp.content.read() == b"" + + +async def test_agents_upload( + hass: HomeAssistant, + hass_client: ClientSessionGenerator, + caplog: pytest.LogCaptureFixture, + mock_api: MagicMock, + snapshot: SnapshotAssertion, +) -> None: + """Test agent upload backup.""" + client = await hass_client() + + with ( + patch( + "homeassistant.components.backup.manager.BackupManager.async_get_backup", + ) as fetch_backup, + patch( + "homeassistant.components.backup.manager.read_backup", + return_value=TEST_AGENT_BACKUP, + ), + patch("pathlib.Path.open") as mocked_open, + ): + mocked_open.return_value.read = Mock(side_effect=[b"test", b""]) + fetch_backup.return_value = TEST_AGENT_BACKUP + resp = await client.post( + f"/api/backup/upload?agent_id=google_drive.{CONFIG_ENTRY_TITLE}", + data={"file": StringIO("test")}, + ) + + assert resp.status == 201 + assert f"Uploading backup {TEST_AGENT_BACKUP.backup_id}" in caplog.text + + mock_api.upload_file.assert_called_once() + assert [tuple(mock_call) for mock_call in mock_api.mock_calls] == snapshot + + +async def test_agents_upload_fail( + hass: HomeAssistant, + hass_client: ClientSessionGenerator, + caplog: pytest.LogCaptureFixture, + mock_api: MagicMock, +) -> None: + """Test agent upload backup fails.""" + mock_api.upload_file = AsyncMock(side_effect=GoogleDriveApiError("some error")) + + client = await hass_client() + + with ( + patch( + "homeassistant.components.backup.manager.BackupManager.async_get_backup", + ) as fetch_backup, + patch( + "homeassistant.components.backup.manager.read_backup", + return_value=TEST_AGENT_BACKUP, + ), + patch("pathlib.Path.open") as mocked_open, + ): + mocked_open.return_value.read = Mock(side_effect=[b"test", b""]) + fetch_backup.return_value = TEST_AGENT_BACKUP + resp = await client.post( + f"/api/backup/upload?agent_id=google_drive.{CONFIG_ENTRY_TITLE}", + data={"file": StringIO("test")}, + ) + await hass.async_block_till_done() + + assert resp.status == 201 + assert "Upload backup error: some error" in caplog.text + + +async def test_agents_delete( + hass: HomeAssistant, + hass_ws_client: WebSocketGenerator, + mock_api: MagicMock, + snapshot: SnapshotAssertion, +) -> None: + """Test agent delete backup.""" + mock_api.list_files = AsyncMock(return_value={"files": [{"id": "backup-file-id"}]}) + mock_api.delete_file = AsyncMock(return_value=None) + + client = await hass_ws_client(hass) + await client.send_json_auto_id( + { + "type": "backup/delete", + "backup_id": TEST_AGENT_BACKUP.backup_id, + } + ) + response = await client.receive_json() + + assert response["success"] + assert response["result"] == {"agent_errors": {}} + + mock_api.delete_file.assert_called_once() + assert [tuple(mock_call) for mock_call in mock_api.mock_calls] == snapshot + + +async def test_agents_delete_fail( + hass: HomeAssistant, + hass_ws_client: WebSocketGenerator, + mock_api: MagicMock, +) -> None: + """Test agent delete backup fails.""" + mock_api.list_files = AsyncMock(return_value={"files": [{"id": "backup-file-id"}]}) + mock_api.delete_file = AsyncMock(side_effect=GoogleDriveApiError("some error")) + + client = await hass_ws_client(hass) + await client.send_json_auto_id( + { + "type": "backup/delete", + "backup_id": TEST_AGENT_BACKUP.backup_id, + } + ) + response = await client.receive_json() + + assert response["success"] + assert response["result"] == { + "agent_errors": { + f"google_drive.{CONFIG_ENTRY_TITLE}": "Failed to delete backup" + } + } + + +async def test_agents_delete_not_found( + hass: HomeAssistant, + hass_ws_client: WebSocketGenerator, + mock_api: MagicMock, +) -> None: + """Test agent delete backup not found.""" + mock_api.list_files = AsyncMock(return_value={"files": []}) + + client = await hass_ws_client(hass) + backup_id = "1234" + + await client.send_json_auto_id( + { + "type": "backup/delete", + "backup_id": backup_id, + } + ) + response = await client.receive_json() + + assert response["success"] + assert response["result"] == {"agent_errors": {}} + + mock_api.delete_file.assert_not_called() diff --git a/tests/components/google_drive/test_config_flow.py b/tests/components/google_drive/test_config_flow.py index 697f88c69a376c..d3c70cbd5eee95 100644 --- a/tests/components/google_drive/test_config_flow.py +++ b/tests/components/google_drive/test_config_flow.py @@ -1,15 +1,16 @@ """Test the Google Drive config flow.""" -from unittest.mock import patch +from unittest.mock import AsyncMock, MagicMock, patch -from aiohttp import ClientError +from google_drive_api.exceptions import GoogleDriveApiError import pytest +from syrupy.assertion import SnapshotAssertion from homeassistant import config_entries from homeassistant.components.google_drive.const import DOMAIN from homeassistant.core import HomeAssistant from homeassistant.data_entry_flow import FlowResultType -from homeassistant.helpers import config_entry_oauth2_flow, instance_id +from homeassistant.helpers import config_entry_oauth2_flow from .conftest import CLIENT_ID @@ -30,6 +31,8 @@ async def test_full_flow( hass: HomeAssistant, hass_client_no_auth: ClientSessionGenerator, aioclient_mock: AiohttpClientMocker, + mock_api: MagicMock, + snapshot: SnapshotAssertion, ) -> None: """Check full flow.""" result = await hass.config_entries.flow.async_init( @@ -56,17 +59,10 @@ async def test_full_flow( assert resp.headers["content-type"] == "text/html; charset=utf-8" # Prepare API responses - aioclient_mock.get( - "https://www.googleapis.com/drive/v3/about", - json={"user": {"emailAddress": USER_EMAIL}}, - ) - aioclient_mock.get( - "https://www.googleapis.com/drive/v3/files", - json={"files": []}, - ) - aioclient_mock.post( - "https://www.googleapis.com/drive/v3/files", - json={"id": FOLDER_ID, "name": FOLDER_NAME}, + mock_api.get_user = AsyncMock(return_value={"user": {"emailAddress": USER_EMAIL}}) + mock_api.list_files = AsyncMock(return_value={"files": []}) + mock_api.create_file = AsyncMock( + return_value={"id": FOLDER_ID, "name": FOLDER_NAME} ) aioclient_mock.post( @@ -86,21 +82,15 @@ async def test_full_flow( assert len(hass.config_entries.async_entries(DOMAIN)) == 1 assert len(mock_setup.mock_calls) == 1 - assert len(aioclient_mock.mock_calls) == 4 - assert aioclient_mock.mock_calls[3][2] == { - "name": "Home Assistant", - "mimeType": "application/vnd.google-apps.folder", - "properties": { - "ha": "root", - "instance_id": await instance_id.async_get(hass), - }, - } - assert aioclient_mock.mock_calls[3][3] == { - "Authorization": "Bearer mock-access-token" - } + assert len(aioclient_mock.mock_calls) == 1 + assert [tuple(mock_call) for mock_call in mock_api.mock_calls] == snapshot assert result.get("type") is FlowResultType.CREATE_ENTRY assert result.get("title") == TITLE + assert result.get("description_placeholders") == { + "folder_name": FOLDER_NAME, + "url": f"https://drive.google.com/drive/folders/{FOLDER_ID}", + } assert "result" in result assert result.get("result").unique_id == USER_EMAIL assert "token" in result.get("result").data @@ -115,6 +105,7 @@ async def test_create_folder_error( hass: HomeAssistant, hass_client_no_auth: ClientSessionGenerator, aioclient_mock: AiohttpClientMocker, + mock_api: MagicMock, ) -> None: """Test case where creating the folder fails.""" result = await hass.config_entries.flow.async_init( @@ -141,18 +132,9 @@ async def test_create_folder_error( assert resp.headers["content-type"] == "text/html; charset=utf-8" # Prepare API responses - aioclient_mock.get( - "https://www.googleapis.com/drive/v3/about", - json={"user": {"emailAddress": USER_EMAIL}}, - ) - aioclient_mock.get( - "https://www.googleapis.com/drive/v3/files", - json={"files": []}, - ) - aioclient_mock.post( - "https://www.googleapis.com/drive/v3/files", - exc=ClientError, - ) + mock_api.get_user = AsyncMock(return_value={"user": {"emailAddress": USER_EMAIL}}) + mock_api.list_files = AsyncMock(return_value={"files": []}) + mock_api.create_file = AsyncMock(side_effect=GoogleDriveApiError("some error")) aioclient_mock.post( GOOGLE_TOKEN_URI, @@ -167,6 +149,7 @@ async def test_create_folder_error( result = await hass.config_entries.flow.async_configure(result["flow_id"]) assert result.get("type") is FlowResultType.ABORT assert result.get("reason") == "create_folder_failure" + assert result.get("description_placeholders") == {"message": "some error"} @pytest.mark.usefixtures("current_request_with_host") @@ -174,6 +157,7 @@ async def test_get_email_error( hass: HomeAssistant, hass_client_no_auth: ClientSessionGenerator, aioclient_mock: AiohttpClientMocker, + mock_api: MagicMock, ) -> None: """Test case where getting the email address fails.""" result = await hass.config_entries.flow.async_init( @@ -200,10 +184,7 @@ async def test_get_email_error( assert resp.headers["content-type"] == "text/html; charset=utf-8" # Prepare API responses - aioclient_mock.get( - "https://www.googleapis.com/drive/v3/about", - exc=ClientError, - ) + mock_api.get_user = AsyncMock(side_effect=GoogleDriveApiError("some error")) aioclient_mock.post( GOOGLE_TOKEN_URI, json={ @@ -216,7 +197,8 @@ async def test_get_email_error( result = await hass.config_entries.flow.async_configure(result["flow_id"]) assert result.get("type") is FlowResultType.ABORT - assert result.get("reason") == "cannot_connect" + assert result.get("reason") == "access_not_configured" + assert result.get("description_placeholders") == {"message": "some error"} @pytest.mark.usefixtures("current_request_with_host") @@ -229,10 +211,10 @@ async def test_get_email_error( ), ( "other.user@domain.com", - "unique_id_mismatch", + "wrong_account", ), ], - ids=["reauth_successful", "unique_id_mismatch"], + ids=["reauth_successful", "wrong_account"], ) async def test_reauth( hass: HomeAssistant, @@ -240,6 +222,7 @@ async def test_reauth( aioclient_mock: AiohttpClientMocker, new_email: str, expected_abort_reason: str, + mock_api: MagicMock, ) -> None: """Test the reauthentication flow.""" @@ -282,10 +265,7 @@ async def test_reauth( assert resp.headers["content-type"] == "text/html; charset=utf-8" # Prepare API responses - aioclient_mock.get( - "https://www.googleapis.com/drive/v3/about", - json={"user": {"emailAddress": new_email}}, - ) + mock_api.get_user = AsyncMock(return_value={"user": {"emailAddress": new_email}}) aioclient_mock.post( GOOGLE_TOKEN_URI, json={ @@ -327,6 +307,7 @@ async def test_already_configured( hass: HomeAssistant, hass_client_no_auth: ClientSessionGenerator, aioclient_mock: AiohttpClientMocker, + mock_api: MagicMock, ) -> None: """Test already configured account.""" config_entry = MockConfigEntry( @@ -364,10 +345,7 @@ async def test_already_configured( assert resp.headers["content-type"] == "text/html; charset=utf-8" # Prepare API responses - aioclient_mock.get( - "https://www.googleapis.com/drive/v3/about", - json={"user": {"emailAddress": USER_EMAIL}}, - ) + mock_api.get_user = AsyncMock(return_value={"user": {"emailAddress": USER_EMAIL}}) aioclient_mock.post( GOOGLE_TOKEN_URI, json={ diff --git a/tests/components/google_drive/test_init.py b/tests/components/google_drive/test_init.py index 09ee0f3a731338..fa52fab84479ad 100644 --- a/tests/components/google_drive/test_init.py +++ b/tests/components/google_drive/test_init.py @@ -4,10 +4,12 @@ import http import time from typing import Any +from unittest.mock import AsyncMock, MagicMock +from google_drive_api.exceptions import GoogleDriveApiError import pytest -from homeassistant.components.google_drive import DOMAIN +from homeassistant.components.google_drive.const import DOMAIN from homeassistant.config_entries import ConfigEntryState from homeassistant.core import HomeAssistant @@ -61,13 +63,12 @@ async def func() -> None: async def test_setup_success( hass: HomeAssistant, setup_integration: ComponentSetup, - aioclient_mock: AiohttpClientMocker, + mock_api: MagicMock, ) -> None: """Test successful setup and unload.""" # Setup looks up existing folder to make sure it still exists - aioclient_mock.get( - "https://www.googleapis.com/drive/v3/files", - json={"files": [{"id": "HA folder ID", "name": "HA folder name"}]}, + mock_api.list_files = AsyncMock( + return_value={"files": [{"id": "HA folder ID", "name": "HA folder name"}]} ) await setup_integration() @@ -84,43 +85,20 @@ async def test_setup_success( assert not hass.services.async_services().get(DOMAIN, {}) -@pytest.mark.parametrize( - ("status", "expected_state"), - [ - ( - http.HTTPStatus.BAD_REQUEST, - ConfigEntryState.SETUP_ERROR, - ), - ( - http.HTTPStatus.NOT_FOUND, - ConfigEntryState.SETUP_ERROR, - ), - ( - http.HTTPStatus.INTERNAL_SERVER_ERROR, - ConfigEntryState.SETUP_RETRY, - ), - ], - ids=["bad_request", "not_found", "transient_failure"], -) async def test_setup_error( hass: HomeAssistant, setup_integration: ComponentSetup, - aioclient_mock: AiohttpClientMocker, - status: http.HTTPStatus, - expected_state: ConfigEntryState, + mock_api: MagicMock, ) -> None: """Test setup error.""" # Simulate failure looking up existing folder - aioclient_mock.get( - "https://www.googleapis.com/drive/v3/files", - status=status, - ) + mock_api.list_files = AsyncMock(side_effect=GoogleDriveApiError("some error")) await setup_integration() entries = hass.config_entries.async_entries(DOMAIN) assert len(entries) == 1 - assert entries[0].state is expected_state + assert entries[0].state is ConfigEntryState.SETUP_RETRY @pytest.mark.parametrize("expires_at", [time.time() - 3600], ids=["expired"]) @@ -128,14 +106,13 @@ async def test_expired_token_refresh_success( hass: HomeAssistant, setup_integration: ComponentSetup, aioclient_mock: AiohttpClientMocker, + mock_api: MagicMock, ) -> None: """Test expired token is refreshed.""" # Setup looks up existing folder to make sure it still exists - aioclient_mock.get( - "https://www.googleapis.com/drive/v3/files", - json={"files": [{"id": "HA folder ID", "name": "HA folder name"}]}, + mock_api.list_files = AsyncMock( + return_value={"files": [{"id": "HA folder ID", "name": "HA folder name"}]} ) - aioclient_mock.post( "https://oauth2.googleapis.com/token", json={ From 1a66485ce8bbc074e5867fd7efdf9f9e771da4ad Mon Sep 17 00:00:00 2001 From: tronikos Date: Sun, 12 Jan 2025 11:32:46 +0000 Subject: [PATCH 22/23] fixes --- homeassistant/components/google_drive/api.py | 6 +- .../components/google_drive/backup.py | 12 +-- .../components/google_drive/strings.json | 4 + tests/components/google_drive/conftest.py | 36 +++++++- .../google_drive/snapshots/test_backup.ambr | 78 +++++++++++++++++ tests/components/google_drive/test_backup.py | 86 +++++++++++-------- .../google_drive/test_config_flow.py | 46 ++++------ tests/components/google_drive/test_init.py | 49 +++++------ 8 files changed, 216 insertions(+), 101 deletions(-) diff --git a/homeassistant/components/google_drive/api.py b/homeassistant/components/google_drive/api.py index 347631dc395f2e..19464192551168 100644 --- a/homeassistant/components/google_drive/api.py +++ b/homeassistant/components/google_drive/api.py @@ -151,7 +151,11 @@ async def async_upload_backup( backup.backup_id, backup_metadata, ) - await self._api.upload_file(backup_metadata, open_stream) + await self._api.upload_file( + backup_metadata, + open_stream, + timeout=ClientTimeout(total=_UPLOAD_AND_DOWNLOAD_TIMEOUT), + ) _LOGGER.debug( "Uploaded backup: %s to: '%s'", backup.backup_id, diff --git a/homeassistant/components/google_drive/backup.py b/homeassistant/components/google_drive/backup.py index fc998a7e641353..b6589fdcd924b9 100644 --- a/homeassistant/components/google_drive/backup.py +++ b/homeassistant/components/google_drive/backup.py @@ -10,6 +10,7 @@ from homeassistant.components.backup import AgentBackup, BackupAgent, BackupAgentError from homeassistant.core import HomeAssistant, callback +from homeassistant.exceptions import HomeAssistantError from homeassistant.helpers.aiohttp_client import ChunkAsyncStreamIterator from . import DATA_BACKUP_AGENT_LISTENERS, GoogleDriveConfigEntry @@ -73,7 +74,7 @@ async def async_upload_backup( """ try: await self._client.async_upload_backup(open_stream, backup) - except (GoogleDriveApiError, TimeoutError) as err: + except (GoogleDriveApiError, HomeAssistantError, TimeoutError) as err: _LOGGER.error("Upload backup error: %s", err) raise BackupAgentError("Failed to upload backup") from err @@ -81,7 +82,8 @@ async def async_list_backups(self, **kwargs: Any) -> list[AgentBackup]: """List backups.""" try: return await self._client.async_list_backups() - except GoogleDriveApiError as err: + except (GoogleDriveApiError, HomeAssistantError, TimeoutError) as err: + _LOGGER.error("List backups error: %s", err) raise BackupAgentError("Failed to list backups") from err async def async_get_backup( @@ -112,8 +114,8 @@ async def async_download_backup( _LOGGER.debug("Downloading file_id: %s", file_id) try: stream = await self._client.async_download(file_id) - except (GoogleDriveApiError, TimeoutError) as err: - _LOGGER.error("Download error: %s", err) + except (GoogleDriveApiError, HomeAssistantError, TimeoutError) as err: + _LOGGER.error("Download backup error: %s", err) raise BackupAgentError("Failed to download backup") from err return ChunkAsyncStreamIterator(stream) _LOGGER.error("Download backup_id: %s not found", backup_id) @@ -135,6 +137,6 @@ async def async_delete_backup( try: await self._client.async_delete(file_id) _LOGGER.debug("Deleted backup_id: %s", backup_id) - except GoogleDriveApiError as err: + except (GoogleDriveApiError, HomeAssistantError, TimeoutError) as err: _LOGGER.error("Delete backup error: %s", err) raise BackupAgentError("Failed to delete backup") from err diff --git a/homeassistant/components/google_drive/strings.json b/homeassistant/components/google_drive/strings.json index 093e2674e4030b..5b06c68b778196 100644 --- a/homeassistant/components/google_drive/strings.json +++ b/homeassistant/components/google_drive/strings.json @@ -3,6 +3,10 @@ "step": { "pick_implementation": { "title": "[%key:common::config_flow::title::oauth2_pick_implementation%]" + }, + "reauth_confirm": { + "title": "[%key:common::config_flow::title::reauth%]", + "description": "The Google Drive integration needs to re-authenticate your account" } }, "abort": { diff --git a/tests/components/google_drive/conftest.py b/tests/components/google_drive/conftest.py index 494e4d88fa8aaa..6921434e0ccbaf 100644 --- a/tests/components/google_drive/conftest.py +++ b/tests/components/google_drive/conftest.py @@ -1,6 +1,7 @@ """PyTest fixtures and test helpers.""" from collections.abc import Generator +import time from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -13,9 +14,13 @@ from homeassistant.core import HomeAssistant from homeassistant.setup import async_setup_component +from tests.common import MockConfigEntry + CLIENT_ID = "1234" CLIENT_SECRET = "5678" HA_UUID = "0a123c" +TEST_USER_EMAIL = "testuser@domain.com" +CONFIG_ENTRY_TITLE = "Google Drive entry title" @pytest.fixture(autouse=True) @@ -44,6 +49,31 @@ def mock_instance_id() -> Generator[AsyncMock]: """Mock instance_id.""" with patch( "homeassistant.components.google_drive.config_flow.instance_id.async_get", - ) as mock_async_get: - mock_async_get.return_value = HA_UUID - yield mock_async_get + return_value=HA_UUID, + ): + yield + + +@pytest.fixture(name="expires_at") +def mock_expires_at() -> int: + """Fixture to set the oauth token expiration time.""" + return time.time() + 3600 + + +@pytest.fixture(name="config_entry") +def mock_config_entry(expires_at: int) -> MockConfigEntry: + """Fixture for MockConfigEntry.""" + return MockConfigEntry( + domain=DOMAIN, + unique_id=TEST_USER_EMAIL, + title=CONFIG_ENTRY_TITLE, + data={ + "auth_implementation": DOMAIN, + "token": { + "access_token": "mock-access-token", + "refresh_token": "mock-refresh-token", + "expires_at": expires_at, + "scope": "https://www.googleapis.com/auth/drive.file", + }, + }, + ) diff --git a/tests/components/google_drive/snapshots/test_backup.ambr b/tests/components/google_drive/snapshots/test_backup.ambr index b535096a3109ff..1c7044f304bafc 100644 --- a/tests/components/google_drive/snapshots/test_backup.ambr +++ b/tests/components/google_drive/snapshots/test_backup.ambr @@ -153,6 +153,84 @@ "CoreBackupReaderWriter.async_receive_backup..open_backup() -> 'AsyncIterator[bytes]'", ), dict({ + 'timeout': dict({ + 'ceil_threshold': 5, + 'connect': None, + 'sock_connect': None, + 'sock_read': None, + 'total': 43200, + }), + }), + ), + ]) +# --- +# name: test_agents_upload_create_folder_if_missing + list([ + tuple( + 'list_files', + tuple( + ), + dict({ + 'params': dict({ + 'fields': 'files(id,name)', + 'q': "properties has { key='ha' and value='root' } and properties has { key='instance_id' and value='0a123c' } and trashed=false", + }), + }), + ), + tuple( + 'list_files', + tuple( + ), + dict({ + 'params': dict({ + 'fields': 'files(id,name)', + 'q': "properties has { key='ha' and value='root' } and properties has { key='instance_id' and value='0a123c' } and trashed=false", + }), + }), + ), + tuple( + 'create_file', + tuple( + ), + dict({ + 'json': dict({ + 'mimeType': 'application/vnd.google-apps.folder', + 'name': 'Home Assistant', + 'properties': dict({ + 'ha': 'root', + 'instance_id': '0a123c', + }), + }), + 'params': dict({ + 'fields': 'id,name', + }), + }), + ), + tuple( + 'upload_file', + tuple( + dict({ + 'description': '{"addons": [{"name": "Test", "slug": "test", "version": "1.0.0"}], "backup_id": "test-backup", "date": "2025-01-01T01:23:45.678Z", "database_included": true, "extra_metadata": {"with_automatic_settings": false}, "folders": [], "homeassistant_included": true, "homeassistant_version": "2024.12.0", "name": "Test", "protected": false, "size": 987}', + 'name': 'Test 2025-01-01T01:23:45.678Z.tar', + 'parents': list([ + 'new folder id', + ]), + 'properties': dict({ + 'backup_id': 'test-backup', + 'ha': 'backup', + 'instance_id': '0a123c', + }), + }), + "CoreBackupReaderWriter.async_receive_backup..open_backup() -> 'AsyncIterator[bytes]'", + ), + dict({ + 'timeout': dict({ + 'ceil_threshold': 5, + 'connect': None, + 'sock_connect': None, + 'sock_read': None, + 'total': 43200, + }), }), ), ]) diff --git a/tests/components/google_drive/test_backup.py b/tests/components/google_drive/test_backup.py index 1f885f20575d89..4baa46a30bcf7d 100644 --- a/tests/components/google_drive/test_backup.py +++ b/tests/components/google_drive/test_backup.py @@ -2,7 +2,6 @@ from io import StringIO import json -import time from typing import Any from unittest.mock import AsyncMock, MagicMock, Mock, patch @@ -11,10 +10,6 @@ import pytest from syrupy.assertion import SnapshotAssertion -from homeassistant.components.application_credentials import ( - ClientCredential, - async_import_client_credential, -) from homeassistant.components.backup import ( DOMAIN as BACKUP_DOMAIN, AddonInfo, @@ -24,13 +19,13 @@ from homeassistant.core import HomeAssistant from homeassistant.setup import async_setup_component +from .conftest import CONFIG_ENTRY_TITLE + from tests.common import MockConfigEntry from tests.test_util.aiohttp import mock_stream from tests.typing import ClientSessionGenerator, WebSocketGenerator -FOLDER_ID = "google-folder-it" -TEST_USER_EMAIL = "testuser@domain.com" -CONFIG_ENTRY_TITLE = "Google Drive entry title" +FOLDER_ID = "google-folder-id" TEST_AGENT_BACKUP = AgentBackup( addons=[AddonInfo(name="Test", slug="test", version="1.0.0")], backup_id="test-backup", @@ -48,25 +43,6 @@ ) -@pytest.fixture(name="config_entry") -def mock_config_entry() -> MockConfigEntry: - """Fixture for MockConfigEntry.""" - return MockConfigEntry( - domain=DOMAIN, - unique_id=TEST_USER_EMAIL, - title=CONFIG_ENTRY_TITLE, - data={ - "auth_implementation": DOMAIN, - "token": { - "access_token": "mock-access-token", - "refresh_token": "mock-refresh-token", - "expires_at": time.time() + 3600, - "scope": "https://www.googleapis.com/auth/drive.file", - }, - }, - ) - - @pytest.fixture(autouse=True) async def setup_integration( hass: HomeAssistant, @@ -76,17 +52,10 @@ async def setup_integration( """Set up Google Drive integration.""" config_entry.add_to_hass(hass) assert await async_setup_component(hass, BACKUP_DOMAIN, {BACKUP_DOMAIN: {}}) - assert await async_setup_component(hass, "application_credentials", {}) - await async_import_client_credential( - hass, - DOMAIN, - ClientCredential("client-id", "client-secret"), - DOMAIN, - ) mock_api.list_files = AsyncMock( return_value={"files": [{"id": "HA folder ID", "name": "HA folder name"}]} ) - assert await async_setup_component(hass, DOMAIN, {}) + await hass.config_entries.async_setup(config_entry.entry_id) await hass.async_block_till_done() @@ -318,6 +287,49 @@ async def test_agents_upload( snapshot: SnapshotAssertion, ) -> None: """Test agent upload backup.""" + mock_api.upload_file = AsyncMock(return_value=None) + + client = await hass_client() + + with ( + patch( + "homeassistant.components.backup.manager.BackupManager.async_get_backup", + ) as fetch_backup, + patch( + "homeassistant.components.backup.manager.read_backup", + return_value=TEST_AGENT_BACKUP, + ), + patch("pathlib.Path.open") as mocked_open, + ): + mocked_open.return_value.read = Mock(side_effect=[b"test", b""]) + fetch_backup.return_value = TEST_AGENT_BACKUP + resp = await client.post( + f"/api/backup/upload?agent_id=google_drive.{CONFIG_ENTRY_TITLE}", + data={"file": StringIO("test")}, + ) + + assert resp.status == 201 + assert f"Uploading backup: {TEST_AGENT_BACKUP.backup_id}" in caplog.text + assert f"Uploaded backup: {TEST_AGENT_BACKUP.backup_id}" in caplog.text + + mock_api.upload_file.assert_called_once() + assert [tuple(mock_call) for mock_call in mock_api.mock_calls] == snapshot + + +async def test_agents_upload_create_folder_if_missing( + hass: HomeAssistant, + hass_client: ClientSessionGenerator, + caplog: pytest.LogCaptureFixture, + mock_api: MagicMock, + snapshot: SnapshotAssertion, +) -> None: + """Test agent upload backup creates folder if missing.""" + mock_api.list_files = AsyncMock(return_value={"files": []}) + mock_api.create_file = AsyncMock( + return_value={"id": "new folder id", "name": "Home Assistant"} + ) + mock_api.upload_file = AsyncMock(return_value=None) + client = await hass_client() with ( @@ -338,8 +350,10 @@ async def test_agents_upload( ) assert resp.status == 201 - assert f"Uploading backup {TEST_AGENT_BACKUP.backup_id}" in caplog.text + assert f"Uploading backup: {TEST_AGENT_BACKUP.backup_id}" in caplog.text + assert f"Uploaded backup: {TEST_AGENT_BACKUP.backup_id}" in caplog.text + mock_api.create_file.assert_called_once() mock_api.upload_file.assert_called_once() assert [tuple(mock_call) for mock_call in mock_api.mock_calls] == snapshot diff --git a/tests/components/google_drive/test_config_flow.py b/tests/components/google_drive/test_config_flow.py index d3c70cbd5eee95..da5718d8fab06f 100644 --- a/tests/components/google_drive/test_config_flow.py +++ b/tests/components/google_drive/test_config_flow.py @@ -12,7 +12,7 @@ from homeassistant.data_entry_flow import FlowResultType from homeassistant.helpers import config_entry_oauth2_flow -from .conftest import CLIENT_ID +from .conftest import CLIENT_ID, TEST_USER_EMAIL from tests.common import MockConfigEntry from tests.test_util.aiohttp import AiohttpClientMocker @@ -20,7 +20,6 @@ GOOGLE_AUTH_URI = "https://accounts.google.com/o/oauth2/v2/auth" GOOGLE_TOKEN_URI = "https://oauth2.googleapis.com/token" -USER_EMAIL = "user@domain.com" FOLDER_ID = "google-folder-id" FOLDER_NAME = "folder name" TITLE = "Google Drive" @@ -59,7 +58,9 @@ async def test_full_flow( assert resp.headers["content-type"] == "text/html; charset=utf-8" # Prepare API responses - mock_api.get_user = AsyncMock(return_value={"user": {"emailAddress": USER_EMAIL}}) + mock_api.get_user = AsyncMock( + return_value={"user": {"emailAddress": TEST_USER_EMAIL}} + ) mock_api.list_files = AsyncMock(return_value={"files": []}) mock_api.create_file = AsyncMock( return_value={"id": FOLDER_ID, "name": FOLDER_NAME} @@ -92,7 +93,7 @@ async def test_full_flow( "url": f"https://drive.google.com/drive/folders/{FOLDER_ID}", } assert "result" in result - assert result.get("result").unique_id == USER_EMAIL + assert result.get("result").unique_id == TEST_USER_EMAIL assert "token" in result.get("result").data assert result.get("result").data["token"].get("access_token") == "mock-access-token" assert ( @@ -132,7 +133,9 @@ async def test_create_folder_error( assert resp.headers["content-type"] == "text/html; charset=utf-8" # Prepare API responses - mock_api.get_user = AsyncMock(return_value={"user": {"emailAddress": USER_EMAIL}}) + mock_api.get_user = AsyncMock( + return_value={"user": {"emailAddress": TEST_USER_EMAIL}} + ) mock_api.list_files = AsyncMock(return_value={"files": []}) mock_api.create_file = AsyncMock(side_effect=GoogleDriveApiError("some error")) @@ -206,7 +209,7 @@ async def test_get_email_error( ("new_email", "expected_abort_reason"), [ ( - USER_EMAIL, + TEST_USER_EMAIL, "reauth_successful", ), ( @@ -219,24 +222,14 @@ async def test_get_email_error( async def test_reauth( hass: HomeAssistant, hass_client_no_auth: ClientSessionGenerator, + config_entry: MockConfigEntry, aioclient_mock: AiohttpClientMocker, new_email: str, expected_abort_reason: str, mock_api: MagicMock, ) -> None: """Test the reauthentication flow.""" - - config_entry = MockConfigEntry( - domain=DOMAIN, - unique_id=USER_EMAIL, - data={ - "token": { - "access_token": "mock-access-token", - }, - }, - ) config_entry.add_to_hass(hass) - config_entry.async_start_reauth(hass) await hass.async_block_till_done() @@ -291,34 +284,25 @@ async def test_reauth( assert result.get("type") is FlowResultType.ABORT assert result.get("reason") == expected_abort_reason - assert config_entry.unique_id == USER_EMAIL + assert config_entry.unique_id == TEST_USER_EMAIL assert "token" in config_entry.data + # Verify access token is refreshed if expected_abort_reason == "reauth_successful": assert config_entry.data["token"].get("access_token") == "updated-access-token" - assert config_entry.data["token"].get("refresh_token") == "mock-refresh-token" else: assert config_entry.data["token"].get("access_token") == "mock-access-token" - assert config_entry.data["token"].get("refresh_token") is None @pytest.mark.usefixtures("current_request_with_host") async def test_already_configured( hass: HomeAssistant, hass_client_no_auth: ClientSessionGenerator, + config_entry: MockConfigEntry, aioclient_mock: AiohttpClientMocker, mock_api: MagicMock, ) -> None: """Test already configured account.""" - config_entry = MockConfigEntry( - domain=DOMAIN, - unique_id=USER_EMAIL, - data={ - "token": { - "access_token": "mock-access-token", - }, - }, - ) config_entry.add_to_hass(hass) result = await hass.config_entries.flow.async_init( @@ -345,7 +329,9 @@ async def test_already_configured( assert resp.headers["content-type"] == "text/html; charset=utf-8" # Prepare API responses - mock_api.get_user = AsyncMock(return_value={"user": {"emailAddress": USER_EMAIL}}) + mock_api.get_user = AsyncMock( + return_value={"user": {"emailAddress": TEST_USER_EMAIL}} + ) aioclient_mock.post( GOOGLE_TOKEN_URI, json={ diff --git a/tests/components/google_drive/test_init.py b/tests/components/google_drive/test_init.py index fa52fab84479ad..cab6e537afb305 100644 --- a/tests/components/google_drive/test_init.py +++ b/tests/components/google_drive/test_init.py @@ -16,35 +16,9 @@ from tests.common import MockConfigEntry from tests.test_util.aiohttp import AiohttpClientMocker -TEST_USER_EMAIL = "testuser@domain.com" - type ComponentSetup = Callable[[], Awaitable[None]] -@pytest.fixture(name="expires_at") -def mock_expires_at() -> int: - """Fixture to set the oauth token expiration time.""" - return time.time() + 3600 - - -@pytest.fixture(name="config_entry") -def mock_config_entry(expires_at: int) -> MockConfigEntry: - """Fixture for MockConfigEntry.""" - return MockConfigEntry( - domain=DOMAIN, - unique_id=TEST_USER_EMAIL, - data={ - "auth_implementation": DOMAIN, - "token": { - "access_token": "mock-access-token", - "refresh_token": "mock-refresh-token", - "expires_at": expires_at, - "scope": "https://www.googleapis.com/auth/drive.file", - }, - }, - ) - - @pytest.fixture(name="setup_integration") async def mock_setup_integration( hass: HomeAssistant, @@ -85,6 +59,29 @@ async def test_setup_success( assert not hass.services.async_services().get(DOMAIN, {}) +async def test_create_folder_if_missing( + hass: HomeAssistant, + setup_integration: ComponentSetup, + mock_api: MagicMock, +) -> None: + """Test folder is created if missing.""" + # Setup looks up existing folder to make sure it still exists + # and creates it if missing + mock_api.list_files = AsyncMock(return_value={"files": []}) + mock_api.create_file = AsyncMock( + return_value={"id": "new folder id", "name": "Home Assistant"} + ) + + await setup_integration() + + entries = hass.config_entries.async_entries(DOMAIN) + assert len(entries) == 1 + assert entries[0].state is ConfigEntryState.LOADED + + mock_api.list_files.assert_called_once() + mock_api.create_file.assert_called_once() + + async def test_setup_error( hass: HomeAssistant, setup_integration: ComponentSetup, From 82627445a50f3ede65ae62a872348e372f8fcfb5 Mon Sep 17 00:00:00 2001 From: tronikos Date: Sun, 12 Jan 2025 12:29:30 +0000 Subject: [PATCH 23/23] strings --- .../components/google_drive/config_flow.py | 10 +++--- .../components/google_drive/strings.json | 15 +++++---- .../google_drive/test_config_flow.py | 33 +++++++++++-------- 3 files changed, 34 insertions(+), 24 deletions(-) diff --git a/homeassistant/components/google_drive/config_flow.py b/homeassistant/components/google_drive/config_flow.py index ea945d466fc0f1..54b3ba7238e0be 100644 --- a/homeassistant/components/google_drive/config_flow.py +++ b/homeassistant/components/google_drive/config_flow.py @@ -4,7 +4,7 @@ from collections.abc import Mapping import logging -from typing import Any +from typing import Any, cast from google_drive_api.exceptions import GoogleDriveApiError @@ -80,10 +80,12 @@ async def async_oauth_create_entry(self, data: dict[str, Any]) -> ConfigFlowResu await self.async_set_unique_id(email_address) if self.source == SOURCE_REAUTH: - self._abort_if_unique_id_mismatch(reason="wrong_account") - return self.async_update_reload_and_abort( - self._get_reauth_entry(), data=data + reauth_entry = self._get_reauth_entry() + self._abort_if_unique_id_mismatch( + reason="wrong_account", + description_placeholders={"email": cast(str, reauth_entry.unique_id)}, ) + return self.async_update_reload_and_abort(reauth_entry, data=data) self._abort_if_unique_id_configured() diff --git a/homeassistant/components/google_drive/strings.json b/homeassistant/components/google_drive/strings.json index 5b06c68b778196..bbee24cf54978d 100644 --- a/homeassistant/components/google_drive/strings.json +++ b/homeassistant/components/google_drive/strings.json @@ -7,24 +7,27 @@ "reauth_confirm": { "title": "[%key:common::config_flow::title::reauth%]", "description": "The Google Drive integration needs to re-authenticate your account" + }, + "auth": { + "title": "Link Google Account" } }, "abort": { - "already_configured": "[%key:common::config_flow::abort::already_configured_account%]", - "already_in_progress": "[%key:common::config_flow::abort::already_in_progress%]", "oauth_error": "[%key:common::config_flow::abort::oauth2_error%]", + "oauth_timeout": "[%key:common::config_flow::abort::oauth2_timeout%]", "missing_configuration": "[%key:common::config_flow::abort::oauth2_missing_configuration%]", + "missing_credentials": "[%key:common::config_flow::abort::oauth2_missing_credentials%]", "authorize_url_timeout": "[%key:common::config_flow::abort::oauth2_authorize_url_timeout%]", "no_url_available": "[%key:common::config_flow::abort::oauth2_no_url_available%]", "user_rejected_authorize": "[%key:common::config_flow::abort::oauth2_user_rejected_authorize%]", - "unknown": "[%key:common::config_flow::error::unknown%]", - "oauth_timeout": "[%key:common::config_flow::abort::oauth2_timeout%]", "oauth_unauthorized": "[%key:common::config_flow::abort::oauth2_unauthorized%]", "oauth_failed": "[%key:common::config_flow::abort::oauth2_failed%]", + "already_configured": "[%key:common::config_flow::abort::already_configured_account%]", + "already_in_progress": "[%key:common::config_flow::abort::already_in_progress%]", + "reauth_successful": "[%key:common::config_flow::abort::reauth_successful%]", "access_not_configured": "Unable to access the Google Drive API:\n\n{message}", "create_folder_failure": "Error while creating Google Drive folder:\n\n{message}", - "wrong_account": "Wrong account: Please authenticate with the right account.", - "reauth_successful": "[%key:common::config_flow::abort::reauth_successful%]" + "wrong_account": "Wrong account: Please authenticate with {email}." }, "create_entry": { "default": "Using [{folder_name}]({url}) folder. Feel free to rename it in Google Drive as you wish." diff --git a/tests/components/google_drive/test_config_flow.py b/tests/components/google_drive/test_config_flow.py index da5718d8fab06f..5dd7590adaabf3 100644 --- a/tests/components/google_drive/test_config_flow.py +++ b/tests/components/google_drive/test_config_flow.py @@ -206,15 +206,21 @@ async def test_get_email_error( @pytest.mark.usefixtures("current_request_with_host") @pytest.mark.parametrize( - ("new_email", "expected_abort_reason"), + ( + "new_email", + "expected_abort_reason", + "expected_placeholders", + "expected_access_token", + "expected_setup_calls", + ), [ - ( - TEST_USER_EMAIL, - "reauth_successful", - ), + (TEST_USER_EMAIL, "reauth_successful", None, "updated-access-token", 1), ( "other.user@domain.com", "wrong_account", + {"email": TEST_USER_EMAIL}, + "mock-access-token", + 0, ), ], ids=["reauth_successful", "wrong_account"], @@ -224,9 +230,12 @@ async def test_reauth( hass_client_no_auth: ClientSessionGenerator, config_entry: MockConfigEntry, aioclient_mock: AiohttpClientMocker, + mock_api: MagicMock, new_email: str, expected_abort_reason: str, - mock_api: MagicMock, + expected_placeholders: dict[str, str] | None, + expected_access_token: str, + expected_setup_calls: int, ) -> None: """Test the reauthentication flow.""" config_entry.add_to_hass(hass) @@ -276,22 +285,18 @@ async def test_reauth( await hass.async_block_till_done() assert len(hass.config_entries.async_entries(DOMAIN)) == 1 - if expected_abort_reason == "reauth_successful": - assert len(mock_setup.mock_calls) == 1 - else: - assert len(mock_setup.mock_calls) == 0 + assert len(mock_setup.mock_calls) == expected_setup_calls assert result.get("type") is FlowResultType.ABORT assert result.get("reason") == expected_abort_reason + assert result.get("description_placeholders") == expected_placeholders assert config_entry.unique_id == TEST_USER_EMAIL assert "token" in config_entry.data # Verify access token is refreshed - if expected_abort_reason == "reauth_successful": - assert config_entry.data["token"].get("access_token") == "updated-access-token" - else: - assert config_entry.data["token"].get("access_token") == "mock-access-token" + assert config_entry.data["token"].get("access_token") == expected_access_token + assert config_entry.data["token"].get("refresh_token") == "mock-refresh-token" @pytest.mark.usefixtures("current_request_with_host")