Migrate cloud settings for all Google entities (#92416)

homeassistant/components/cloud/google_config.py

@@ -1,5 +1,6 @@
 """Google config for Cloud."""
 import asyncio
+from contextlib import suppress
 from http import HTTPStatus
 import logging
 from typing import Any
@@ -11,8 +12,10 @@ from homeassistant.components.binary_sensor import BinarySensorDeviceClass
 from homeassistant.components.google_assistant import DOMAIN as GOOGLE_DOMAIN
 from homeassistant.components.google_assistant.helpers import AbstractConfig
 from homeassistant.components.homeassistant.exposed_entities import (
+    async_expose_entity,
     async_get_entity_settings,
     async_listen_entity_updates,
+    async_set_assistant_option,
     async_should_expose,
 )
 from homeassistant.components.sensor import SensorDeviceClass
@@ -173,34 +176,67 @@ class CloudGoogleConfig(AbstractConfig):
             # Don't migrate if there's a YAML config
             return
 
-        entity_registry = er.async_get(self.hass)
+        for state in self.hass.states.async_all():
-
+            entity_id = state.entity_id
-        for entity_id, entry in entity_registry.entities.items():
+            with suppress(HomeAssistantError):
-            if CLOUD_GOOGLE in entry.options:
+                entity_settings = async_get_entity_settings(self.hass, entity_id)
-                continue
+                if CLOUD_GOOGLE in entity_settings:
-            options = {"should_expose": self._should_expose_legacy(entity_id)}
+                    continue
-            if _2fa_disabled := (self._2fa_disabled_legacy(entity_id) is not None):
+            async_expose_entity(
-                options[PREF_DISABLE_2FA] = _2fa_disabled
+                self.hass,
-            entity_registry.async_update_entity_options(
+                CLOUD_GOOGLE,
-                entity_id, CLOUD_GOOGLE, options
+                entity_id,
+                self._should_expose_legacy(entity_id),
             )
+            if _2fa_disabled := (self._2fa_disabled_legacy(entity_id) is not None):
+                async_set_assistant_option(
+                    self.hass,
+                    CLOUD_GOOGLE,
+                    entity_id,
+                    PREF_DISABLE_2FA,
+                    _2fa_disabled,
+                )
+        for entity_id in self._prefs.google_entity_configs:
+            with suppress(HomeAssistantError):
+                entity_settings = async_get_entity_settings(self.hass, entity_id)
+                if CLOUD_GOOGLE in entity_settings:
+                    continue
+            async_expose_entity(
+                self.hass,
+                CLOUD_GOOGLE,
+                entity_id,
+                self._should_expose_legacy(entity_id),
+            )
+            if _2fa_disabled := (self._2fa_disabled_legacy(entity_id) is not None):
+                async_set_assistant_option(
+                    self.hass,
+                    CLOUD_GOOGLE,
+                    entity_id,
+                    PREF_DISABLE_2FA,
+                    _2fa_disabled,
+                )
 
     async def async_initialize(self):
         """Perform async initialization of config."""
         await super().async_initialize()
 
-        if self._prefs.google_settings_version != GOOGLE_SETTINGS_VERSION:
+        async def on_hass_started(hass: HomeAssistant) -> None:
-            if self._prefs.google_settings_version < 2:
+            if self._prefs.google_settings_version != GOOGLE_SETTINGS_VERSION:
-                self._migrate_google_entity_settings_v1()
+                if self._prefs.google_settings_version < 2:
-            await self._prefs.async_update(
+                    self._migrate_google_entity_settings_v1()
-                google_settings_version=GOOGLE_SETTINGS_VERSION
+                await self._prefs.async_update(
+                    google_settings_version=GOOGLE_SETTINGS_VERSION
+                )
+            async_listen_entity_updates(
+                self.hass, CLOUD_GOOGLE, self._async_exposed_entities_updated
             )
 
-        async def hass_started(hass):
+        async def on_hass_start(hass: HomeAssistant) -> None:
             if self.enabled and GOOGLE_DOMAIN not in self.hass.config.components:
                 await async_setup_component(self.hass, GOOGLE_DOMAIN, {})
 
-        start.async_at_start(self.hass, hass_started)
+        start.async_at_start(self.hass, on_hass_start)
+        start.async_at_started(self.hass, on_hass_started)
 
         # Remove any stored user agent id that is not ours
         remove_agent_user_ids = []
@@ -212,9 +248,6 @@ class CloudGoogleConfig(AbstractConfig):
             await self.async_disconnect_agent_user(agent_user_id)
 
         self._prefs.async_listen_updates(self._async_prefs_updated)
-        async_listen_entity_updates(
-            self.hass, CLOUD_GOOGLE, self._async_exposed_entities_updated
-        )
         self.hass.bus.async_listen(
             er.EVENT_ENTITY_REGISTRY_UPDATED,
             self._handle_entity_registry_updated,

tests/components/cloud/test_google_config.py

@@ -21,9 +21,12 @@ from homeassistant.components.homeassistant.exposed_entities import (
     async_expose_entity,
     async_get_entity_settings,
 )
-from homeassistant.const import EVENT_HOMEASSISTANT_STARTED, EntityCategory
+from homeassistant.const import (
+    EVENT_HOMEASSISTANT_START,
+    EVENT_HOMEASSISTANT_STARTED,
+    EntityCategory,
+)
 from homeassistant.core import CoreState, HomeAssistant, State
-from homeassistant.exceptions import HomeAssistantError
 from homeassistant.helpers import device_registry as dr, entity_registry as er
 from homeassistant.setup import async_setup_component
 from homeassistant.util.dt import utcnow
@@ -145,6 +148,8 @@ async def test_google_update_expose_trigger_sync(
             Mock(claims={"cognito:username": "abcdefghjkl"}),
         )
         await config.async_initialize()
+        hass.bus.async_fire(EVENT_HOMEASSISTANT_STARTED)
+        await hass.async_block_till_done()
         await config.async_connect_agent_user("mock-user-id")
 
         with patch.object(config, "async_sync_entities") as mock_sync, patch.object(
@@ -484,8 +489,10 @@ async def test_google_config_migrate_expose_entity_prefs(
     entity_registry: er.EntityRegistry,
 ) -> None:
     """Test migrating Google entity config."""
+    hass.state = CoreState.starting
 
     assert await async_setup_component(hass, "homeassistant", {})
+    hass.states.async_set("light.state_only", "on")
     entity_exposed = entity_registry.async_get_or_create(
         "light",
         "test",
@@ -538,7 +545,11 @@ async def test_google_config_migrate_expose_entity_prefs(
     expose_entity(hass, entity_migrated.entity_id, False)
 
     cloud_prefs._prefs[PREF_GOOGLE_ENTITY_CONFIGS]["light.unknown"] = {
-        PREF_SHOULD_EXPOSE: True
+        PREF_SHOULD_EXPOSE: True,
+        PREF_DISABLE_2FA: True,
+    }
+    cloud_prefs._prefs[PREF_GOOGLE_ENTITY_CONFIGS]["light.state_only"] = {
+        PREF_SHOULD_EXPOSE: False
     }
     cloud_prefs._prefs[PREF_GOOGLE_ENTITY_CONFIGS][entity_exposed.entity_id] = {
         PREF_SHOULD_EXPOSE: True
@@ -554,9 +565,17 @@ async def test_google_config_migrate_expose_entity_prefs(
         hass, GACTIONS_SCHEMA({}), "mock-user-id", cloud_prefs, Mock(is_logged_in=False)
     )
     await conf.async_initialize()
+    hass.bus.async_fire(EVENT_HOMEASSISTANT_START)
+    await hass.async_block_till_done()
+    hass.bus.async_fire(EVENT_HOMEASSISTANT_STARTED)
+    await hass.async_block_till_done()
 
-    with pytest.raises(HomeAssistantError):
+    assert async_get_entity_settings(hass, "light.unknown") == {
-        async_get_entity_settings(hass, "light.unknown")
+        "cloud.google_assistant": {"disable_2fa": True, "should_expose": True}
+    }
+    assert async_get_entity_settings(hass, "light.state_only") == {
+        "cloud.google_assistant": {"should_expose": False}
+    }
     assert async_get_entity_settings(hass, entity_exposed.entity_id) == {
         "cloud.google_assistant": {"should_expose": True}
     }
@@ -583,6 +602,7 @@ async def test_google_config_migrate_expose_entity_prefs_default_none(
     entity_registry: er.EntityRegistry,
 ) -> None:
     """Test migrating Google entity config."""
+    hass.state = CoreState.starting
 
     assert await async_setup_component(hass, "homeassistant", {})
     entity_default = entity_registry.async_get_or_create(
@@ -603,6 +623,10 @@ async def test_google_config_migrate_expose_entity_prefs_default_none(
         hass, GACTIONS_SCHEMA({}), "mock-user-id", cloud_prefs, Mock(is_logged_in=False)
     )
     await conf.async_initialize()
+    hass.bus.async_fire(EVENT_HOMEASSISTANT_START)
+    await hass.async_block_till_done()
+    hass.bus.async_fire(EVENT_HOMEASSISTANT_STARTED)
+    await hass.async_block_till_done()
 
     assert async_get_entity_settings(hass, entity_default.entity_id) == {
         "cloud.google_assistant": {"should_expose": True}
@@ -615,6 +639,7 @@ async def test_google_config_migrate_expose_entity_prefs_default(
     entity_registry: er.EntityRegistry,
 ) -> None:
     """Test migrating Google entity config."""
+    hass.state = CoreState.starting
 
     assert await async_setup_component(hass, "homeassistant", {})
 
@@ -680,6 +705,10 @@ async def test_google_config_migrate_expose_entity_prefs_default(
         hass, GACTIONS_SCHEMA({}), "mock-user-id", cloud_prefs, Mock(is_logged_in=False)
     )
     await conf.async_initialize()
+    hass.bus.async_fire(EVENT_HOMEASSISTANT_START)
+    await hass.async_block_till_done()
+    hass.bus.async_fire(EVENT_HOMEASSISTANT_STARTED)
+    await hass.async_block_till_done()
 
     assert async_get_entity_settings(hass, binary_sensor_supported.entity_id) == {
         "cloud.google_assistant": {"should_expose": True}