Store runtime data inside the config entry in Google Sheets (#119438)

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(
+        service = Client(Credentials(entry.data[CONF_TOKEN][CONF_ACCESS_TOKEN]))  # type: ignore[no-untyped-call]
-            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)):
+        await entry.runtime_data.async_ensure_token_valid()
-            raise ValueError(f"Config entry not loaded: {call.data[DATA_CONFIG_ENTRY]}")
-        await session.async_ensure_token_valid()
         await hass.async_add_executor_job(_append_to_sheet, call, entry)
 
     hass.services.async_register(

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:

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",