Use runtime_data in google_assistant (#144332)

homeassistant/components/google_assistant/__init__.py

@@ -95,6 +95,8 @@ CONFIG_SCHEMA = vol.Schema(
     {vol.Optional(DOMAIN): GOOGLE_ASSISTANT_SCHEMA}, extra=vol.ALLOW_EXTRA
 )
 
+type GoogleConfigEntry = ConfigEntry[GoogleConfig]
+
 
 async def async_setup(hass: HomeAssistant, yaml_config: ConfigType) -> bool:
     """Activate Google Actions component."""
@@ -115,7 +117,7 @@ async def async_setup(hass: HomeAssistant, yaml_config: ConfigType) -> bool:
     return True
 
 
-async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+async def async_setup_entry(hass: HomeAssistant, entry: GoogleConfigEntry) -> bool:
     """Set up from a config entry."""
 
     config: ConfigType = {**hass.data[DOMAIN][DATA_CONFIG]}
@@ -141,7 +143,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     google_config = GoogleConfig(hass, config)
     await google_config.async_initialize()
 
-    hass.data[DOMAIN][entry.entry_id] = google_config
+    entry.runtime_data = google_config
 
     hass.http.register_view(GoogleAssistantView(google_config))
 

homeassistant/components/google_assistant/button.py

@@ -2,7 +2,6 @@
 
 from __future__ import annotations
 
-from homeassistant import config_entries
 from homeassistant.components.button import ButtonEntity
 from homeassistant.const import EntityCategory
 from homeassistant.core import HomeAssistant
@@ -11,18 +10,19 @@ from homeassistant.helpers.device_registry import DeviceInfo
 from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
 from homeassistant.helpers.typing import ConfigType
 
+from . import GoogleConfigEntry
 from .const import CONF_PROJECT_ID, CONF_SERVICE_ACCOUNT, DATA_CONFIG, DOMAIN
 from .http import GoogleConfig
 
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    config_entry: config_entries.ConfigEntry,
+    config_entry: GoogleConfigEntry,
     async_add_entities: AddConfigEntryEntitiesCallback,
 ) -> None:
     """Set up the platform."""
     yaml_config: ConfigType = hass.data[DOMAIN][DATA_CONFIG]
-    google_config: GoogleConfig = hass.data[DOMAIN][config_entry.entry_id]
+    google_config = config_entry.runtime_data
 
     entities = []
 

homeassistant/components/google_assistant/diagnostics.py

@@ -5,13 +5,12 @@ from __future__ import annotations
 from typing import Any
 
 from homeassistant.components.diagnostics import REDACTED, async_redact_data
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_API_KEY
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.typing import ConfigType
 
+from . import GoogleConfigEntry
 from .const import CONF_SECURE_DEVICES_PIN, CONF_SERVICE_ACCOUNT, DATA_CONFIG, DOMAIN
-from .http import GoogleConfig
 from .smart_home import (
     async_devices_query_response,
     async_devices_sync_response,
@@ -29,12 +28,11 @@ TO_REDACT = [
 
 
 async def async_get_config_entry_diagnostics(
-    hass: HomeAssistant, entry: ConfigEntry
+    hass: HomeAssistant, entry: GoogleConfigEntry
 ) -> dict[str, Any]:
     """Return diagnostic information."""
-    data = hass.data[DOMAIN]
-    config: GoogleConfig = data[entry.entry_id]
-    yaml_config: ConfigType = data[DATA_CONFIG]
+    config = entry.runtime_data
+    yaml_config: ConfigType = hass.data[DOMAIN][DATA_CONFIG]
     devices = await async_devices_sync_response(hass, config, REDACTED)
     sync = create_sync_response(REDACTED, devices)
     query = await async_devices_query_response(hass, config, devices)

tests/components/google_assistant/test_button.py

@@ -29,7 +29,7 @@ async def test_sync_button(hass: HomeAssistant, hass_owner_user: MockUser) -> No
     assert state
 
     config_entry = hass.config_entries.async_entries("google_assistant")[0]
-    google_config: ga.GoogleConfig = hass.data[ga.DOMAIN][config_entry.entry_id]
+    google_config: ga.GoogleConfig = config_entry.runtime_data
 
     with patch.object(google_config, "async_sync_entities") as mock_sync_entities:
         mock_sync_entities.return_value = 200