Store runtime data inside the config entry in Google Mail (#119439)

homeassistant/components/google_mail/__init__.py

@@ -16,6 +16,8 @@ from .api import AsyncConfigEntryAuth
 from .const import DATA_AUTH, DATA_HASS_CONFIG, DOMAIN
 from .services import async_setup_services
 
+type GoogleMailConfigEntry = ConfigEntry[AsyncConfigEntryAuth]
+
 PLATFORMS = [Platform.NOTIFY, Platform.SENSOR]
 
 CONFIG_SCHEMA = cv.config_entry_only_config_schema(DOMAIN)
@@ -28,13 +30,13 @@ async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
     return True
 
 
-async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+async def async_setup_entry(hass: HomeAssistant, entry: GoogleMailConfigEntry) -> bool:
     """Set up Google Mail 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()
-    hass.data[DOMAIN][entry.entry_id] = auth
+    entry.runtime_data = auth
 
     hass.async_create_task(
         discovery.async_load_platform(
@@ -55,10 +57,8 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     return True
 
 
-async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+async def async_unload_entry(hass: HomeAssistant, entry: GoogleMailConfigEntry) -> bool:
     """Unload a config entry."""
-    if unload_ok := await hass.config_entries.async_unload_platforms(entry, PLATFORMS):
-        hass.data[DOMAIN].pop(entry.entry_id)
     loaded_entries = [
         entry
         for entry in hass.config_entries.async_entries(DOMAIN)
@@ -68,4 +68,4 @@ async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
         for service_name in hass.services.async_services_for_domain(DOMAIN):
             hass.services.async_remove(DOMAIN, service_name)
 
-    return unload_ok
+    return await hass.config_entries.async_unload_platforms(entry, PLATFORMS)

homeassistant/components/google_mail/config_flow.py

@@ -9,10 +9,11 @@ from typing import Any, cast
 from google.oauth2.credentials import Credentials
 from googleapiclient.discovery import build
 
-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 GoogleMailConfigEntry
 from .const import DEFAULT_ACCESS, DOMAIN
 
 
@@ -23,7 +24,7 @@ class OAuth2FlowHandler(
 
     DOMAIN = DOMAIN
 
-    reauth_entry: ConfigEntry | None = None
+    reauth_entry: GoogleMailConfigEntry | None = None
 
     @property
     def logger(self) -> logging.Logger:

homeassistant/components/google_mail/sensor.py

@@ -11,11 +11,10 @@ from homeassistant.components.sensor import (
     SensorEntity,
     SensorEntityDescription,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from .const import DOMAIN
+from . import GoogleMailConfigEntry
 from .entity import GoogleMailEntity
 
 SCAN_INTERVAL = timedelta(minutes=15)
@@ -28,12 +27,12 @@ SENSOR_TYPE = SensorEntityDescription(
 
 
 async def async_setup_entry(
-    hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
+    hass: HomeAssistant,
+    entry: GoogleMailConfigEntry,
+    async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up the Google Mail sensor."""
-    async_add_entities(
-        [GoogleMailSensor(hass.data[DOMAIN][entry.entry_id], SENSOR_TYPE)], True
-    )
+    async_add_entities([GoogleMailSensor(entry.runtime_data, SENSOR_TYPE)], True)
 
 
 class GoogleMailSensor(GoogleMailEntity, SensorEntity):

homeassistant/components/google_mail/services.py

@@ -3,16 +3,15 @@
 from __future__ import annotations
 
 from datetime import datetime, timedelta
+from typing import TYPE_CHECKING
 
 from googleapiclient.http import HttpRequest
 import voluptuous as vol
 
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant, ServiceCall
 from homeassistant.helpers import config_validation as cv
 from homeassistant.helpers.service import async_extract_config_entry_ids
 
-from .api import AsyncConfigEntryAuth
 from .const import (
     ATTR_ENABLED,
     ATTR_END,
@@ -26,6 +25,9 @@ from .const import (
     DOMAIN,
 )
 
+if TYPE_CHECKING:
+    from . import GoogleMailConfigEntry
+
 SERVICE_SET_VACATION = "set_vacation"
 
 SERVICE_VACATION_SCHEMA = vol.All(
@@ -47,7 +49,9 @@ SERVICE_VACATION_SCHEMA = vol.All(
 async def async_setup_services(hass: HomeAssistant) -> None:
     """Set up services for Google Mail integration."""
 
-    async def extract_gmail_config_entries(call: ServiceCall) -> list[ConfigEntry]:
+    async def extract_gmail_config_entries(
+        call: ServiceCall,
+    ) -> list[GoogleMailConfigEntry]:
         return [
             entry
             for entry_id in await async_extract_config_entry_ids(hass, call)
@@ -57,10 +61,11 @@ async def async_setup_services(hass: HomeAssistant) -> None:
 
     async def gmail_service(call: ServiceCall) -> None:
         """Call Google Mail service."""
-        auth: AsyncConfigEntryAuth
         for entry in await extract_gmail_config_entries(call):
-            if not (auth := hass.data[DOMAIN].get(entry.entry_id)):
-                raise ValueError(f"Config entry not loaded: {entry.entry_id}")
+            try:
+                auth = entry.runtime_data
+            except AttributeError as ex:
+                raise ValueError(f"Config entry not loaded: {entry.entry_id}") from ex
             service = await auth.get_resource()
 
             _settings = {