From a9d8bc989ec37d9047be2cbffbd0697bd43eb5fc Mon Sep 17 00:00:00 2001 From: Erik Montnemery Date: Wed, 3 May 2023 18:56:48 +0200 Subject: [PATCH] Migrate cloud settings for all Google entities (#92416) --- .../components/cloud/google_config.py | 73 ++++++++++++++----- tests/components/cloud/test_google_config.py | 39 ++++++++-- 2 files changed, 87 insertions(+), 25 deletions(-) diff --git a/homeassistant/components/cloud/google_config.py b/homeassistant/components/cloud/google_config.py index 03dd27c7c38..dae1c00a33f 100644 --- a/homeassistant/components/cloud/google_config.py +++ b/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 entity_id, entry in entity_registry.entities.items(): - if CLOUD_GOOGLE in entry.options: - continue - options = {"should_expose": self._should_expose_legacy(entity_id)} - if _2fa_disabled := (self._2fa_disabled_legacy(entity_id) is not None): - options[PREF_DISABLE_2FA] = _2fa_disabled - entity_registry.async_update_entity_options( - entity_id, CLOUD_GOOGLE, options + for state in self.hass.states.async_all(): + entity_id = state.entity_id + 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, + ) + 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: - if self._prefs.google_settings_version < 2: - self._migrate_google_entity_settings_v1() - await self._prefs.async_update( - google_settings_version=GOOGLE_SETTINGS_VERSION + async def on_hass_started(hass: HomeAssistant) -> None: + if self._prefs.google_settings_version != GOOGLE_SETTINGS_VERSION: + if self._prefs.google_settings_version < 2: + self._migrate_google_entity_settings_v1() + 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, diff --git a/tests/components/cloud/test_google_config.py b/tests/components/cloud/test_google_config.py index 0738bc6b029..0fa37ed9987 100644 --- a/tests/components/cloud/test_google_config.py +++ b/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): - async_get_entity_settings(hass, "light.unknown") + assert 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}