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

This commit is contained in:
Robert Hillis 2024-06-12 06:48:55 -04:00 committed by GitHub
parent ade936e6d5
commit 35b13e355b
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 19 additions and 17 deletions

View file

@ -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(

View file

@ -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:

View file

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