diff --git a/homeassistant/components/google_sheets/__init__.py b/homeassistant/components/google_sheets/__init__.py index 713a801257d..fc104cc5c22 100644 --- a/homeassistant/components/google_sheets/__init__.py +++ b/homeassistant/components/google_sheets/__init__.py @@ -29,6 +29,8 @@ from homeassistant.helpers.selector import ConfigEntrySelector from .const import DEFAULT_ACCESS, DOMAIN +type GoogleSheetsConfigEntry = ConfigEntry[OAuth2Session] + DATA = "data" DATA_CONFIG_ENTRY = "config_entry" WORKSHEET = "worksheet" @@ -44,7 +46,9 @@ SHEET_SERVICE_SCHEMA = vol.All( ) -async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: +async def async_setup_entry( + hass: HomeAssistant, entry: GoogleSheetsConfigEntry +) -> bool: """Set up Google Sheets from a config entry.""" implementation = await async_get_config_entry_implementation(hass, entry) session = OAuth2Session(hass, entry, implementation) @@ -61,21 +65,22 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: if not async_entry_has_scopes(hass, entry): raise ConfigEntryAuthFailed("Required scopes are not present, reauth required") - hass.data.setdefault(DOMAIN, {})[entry.entry_id] = session + entry.runtime_data = session await async_setup_service(hass) return True -def async_entry_has_scopes(hass: HomeAssistant, entry: ConfigEntry) -> bool: +def async_entry_has_scopes(hass: HomeAssistant, entry: GoogleSheetsConfigEntry) -> bool: """Verify that the config entry desired scope is present in the oauth token.""" return DEFAULT_ACCESS in entry.data.get(CONF_TOKEN, {}).get("scope", "").split(" ") -async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: +async def async_unload_entry( + hass: HomeAssistant, entry: GoogleSheetsConfigEntry +) -> bool: """Unload a config entry.""" - hass.data[DOMAIN].pop(entry.entry_id) loaded_entries = [ entry for entry in hass.config_entries.async_entries(DOMAIN) @@ -91,11 +96,9 @@ async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: async def async_setup_service(hass: HomeAssistant) -> None: """Add the services for Google Sheets.""" - def _append_to_sheet(call: ServiceCall, entry: ConfigEntry) -> None: + def _append_to_sheet(call: ServiceCall, entry: GoogleSheetsConfigEntry) -> None: """Run append in the executor.""" - service = Client( - Credentials(entry.data[CONF_TOKEN][CONF_ACCESS_TOKEN]) # type: ignore[no-untyped-call] - ) + service = Client(Credentials(entry.data[CONF_TOKEN][CONF_ACCESS_TOKEN])) # type: ignore[no-untyped-call] try: sheet = service.open_by_key(entry.unique_id) except RefreshError: @@ -117,14 +120,12 @@ async def async_setup_service(hass: HomeAssistant) -> None: async def append_to_sheet(call: ServiceCall) -> None: """Append new line of data to a Google Sheets document.""" - entry: ConfigEntry | None = hass.config_entries.async_get_entry( + entry: GoogleSheetsConfigEntry | None = hass.config_entries.async_get_entry( call.data[DATA_CONFIG_ENTRY] ) - if not entry: + if not entry or not hasattr(entry, "runtime_data"): raise ValueError(f"Invalid config entry: {call.data[DATA_CONFIG_ENTRY]}") - if not (session := hass.data[DOMAIN].get(entry.entry_id)): - raise ValueError(f"Config entry not loaded: {call.data[DATA_CONFIG_ENTRY]}") - await session.async_ensure_token_valid() + await entry.runtime_data.async_ensure_token_valid() await hass.async_add_executor_job(_append_to_sheet, call, entry) hass.services.async_register( diff --git a/homeassistant/components/google_sheets/config_flow.py b/homeassistant/components/google_sheets/config_flow.py index ab0c084c317..4008d42f52d 100644 --- a/homeassistant/components/google_sheets/config_flow.py +++ b/homeassistant/components/google_sheets/config_flow.py @@ -9,10 +9,11 @@ from typing import Any from google.oauth2.credentials import Credentials from gspread import Client, GSpreadException -from homeassistant.config_entries import ConfigEntry, ConfigFlowResult +from homeassistant.config_entries import ConfigFlowResult from homeassistant.const import CONF_ACCESS_TOKEN, CONF_TOKEN from homeassistant.helpers import config_entry_oauth2_flow +from . import GoogleSheetsConfigEntry from .const import DEFAULT_ACCESS, DEFAULT_NAME, DOMAIN _LOGGER = logging.getLogger(__name__) @@ -25,7 +26,7 @@ class OAuth2FlowHandler( DOMAIN = DOMAIN - reauth_entry: ConfigEntry | None = None + reauth_entry: GoogleSheetsConfigEntry | None = None @property def logger(self) -> logging.Logger: diff --git a/tests/components/google_sheets/test_init.py b/tests/components/google_sheets/test_init.py index 0842debc38d..014e89349e2 100644 --- a/tests/components/google_sheets/test_init.py +++ b/tests/components/google_sheets/test_init.py @@ -294,7 +294,7 @@ async def test_append_sheet_invalid_config_entry( await hass.async_block_till_done() assert config_entry2.state is ConfigEntryState.NOT_LOADED - with pytest.raises(ValueError, match="Config entry not loaded"): + with pytest.raises(ValueError, match="Invalid config entry"): await hass.services.async_call( DOMAIN, "append_sheet",