Skip to content

Commit

Permalink
Migrate cloud settings for all Google entities (home-assistant#92416)
Browse files Browse the repository at this point in the history
  • Loading branch information
emontnemery authored and dknowles2 committed May 4, 2023
1 parent 80444ce commit 61c5867
Show file tree
Hide file tree
Showing 2 changed files with 86 additions and 24 deletions.
71 changes: 52 additions & 19 deletions 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
Expand All @@ -11,8 +12,10 @@
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
Expand Down Expand Up @@ -173,34 +176,67 @@ def _migrate_google_entity_settings_v1(self):
# 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)}
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):
options[PREF_DISABLE_2FA] = _2fa_disabled
entity_registry.async_update_entity_options(
entity_id, CLOUD_GOOGLE, options
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 = []
Expand All @@ -212,9 +248,6 @@ async def hass_started(hass):
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,
Expand Down
39 changes: 34 additions & 5 deletions tests/components/cloud/test_google_config.py
Expand Up @@ -21,9 +21,12 @@
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
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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",
Expand Down Expand Up @@ -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
Expand All @@ -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}
}
Expand All @@ -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(
Expand All @@ -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}
Expand All @@ -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", {})

Expand Down Expand Up @@ -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}
Expand Down

0 comments on commit 61c5867

Please sign in to comment.