From 31de1b17e86210535fcaca5a680753fca13fae21 Mon Sep 17 00:00:00 2001 From: Erik Montnemery Date: Wed, 3 May 2023 15:55:38 +0200 Subject: [PATCH] Allow setting google disable 2fa flag on any entity (#92403) * Allow setting google disable 2fa flag on any entity * Fix test * Include disable_2fa flag in cloud/google_assistant/entities/get --- .../components/cloud/google_config.py | 11 +-- homeassistant/components/cloud/http_api.py | 31 ++++----- .../homeassistant/exposed_entities.py | 67 +++++++++++-------- tests/components/cloud/test_alexa_config.py | 4 +- tests/components/cloud/test_client.py | 5 +- tests/components/cloud/test_google_config.py | 6 +- tests/components/cloud/test_http_api.py | 37 ++++++++-- tests/components/conversation/__init__.py | 4 +- .../google_assistant/test_smart_home.py | 2 +- .../homeassistant/test_exposed_entities.py | 28 ++++---- 10 files changed, 113 insertions(+), 82 deletions(-) diff --git a/homeassistant/components/cloud/google_config.py b/homeassistant/components/cloud/google_config.py index a5700789112..03dd27c7c38 100644 --- a/homeassistant/components/cloud/google_config.py +++ b/homeassistant/components/cloud/google_config.py @@ -11,6 +11,7 @@ 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_get_entity_settings, async_listen_entity_updates, async_should_expose, ) @@ -23,6 +24,7 @@ from homeassistant.core import ( callback, split_entity_id, ) +from homeassistant.exceptions import HomeAssistantError from homeassistant.helpers import device_registry as dr, entity_registry as er, start from homeassistant.helpers.entity import get_device_class from homeassistant.setup import async_setup_component @@ -289,14 +291,13 @@ class CloudGoogleConfig(AbstractConfig): def should_2fa(self, state): """If an entity should be checked for 2FA.""" - entity_registry = er.async_get(self.hass) - - registry_entry = entity_registry.async_get(state.entity_id) - if not registry_entry: + try: + settings = async_get_entity_settings(self.hass, state.entity_id) + except HomeAssistantError: # Handle the entity has been removed return False - assistant_options = registry_entry.options.get(CLOUD_GOOGLE, {}) + assistant_options = settings.get(CLOUD_GOOGLE, {}) return not assistant_options.get(PREF_DISABLE_2FA, DEFAULT_DISABLE_2FA) async def async_report_state(self, message, agent_user_id: str): diff --git a/homeassistant/components/cloud/http_api.py b/homeassistant/components/cloud/http_api.py index 8bc31f7b862..391726cb900 100644 --- a/homeassistant/components/cloud/http_api.py +++ b/homeassistant/components/cloud/http_api.py @@ -1,6 +1,7 @@ """The HTTP api to control the cloud integration.""" import asyncio from collections.abc import Mapping +from contextlib import suppress import dataclasses from functools import wraps from http import HTTPStatus @@ -21,10 +22,12 @@ from homeassistant.components.alexa import ( errors as alexa_errors, ) from homeassistant.components.google_assistant import helpers as google_helpers +from homeassistant.components.homeassistant import exposed_entities from homeassistant.components.http import HomeAssistantView from homeassistant.components.http.data_validator import RequestDataValidator from homeassistant.const import CLOUD_NEVER_EXPOSED_ENTITIES from homeassistant.core import HomeAssistant +from homeassistant.exceptions import HomeAssistantError from homeassistant.helpers import entity_registry as er from homeassistant.helpers.aiohttp_client import async_get_clientsession from homeassistant.util.location import async_detect_location_info @@ -587,10 +590,16 @@ async def google_assistant_get( ) return + assistant_options: Mapping[str, Any] = {} + with suppress(HomeAssistantError, KeyError): + settings = exposed_entities.async_get_entity_settings(hass, entity_id) + assistant_options = settings[CLOUD_GOOGLE] + result = { "entity_id": entity.entity_id, "traits": [trait.name for trait in entity.traits()], "might_2fa": entity.might_2fa_traits(), + PREF_DISABLE_2FA: assistant_options.get(PREF_DISABLE_2FA), } connection.send_result(msg["id"], result) @@ -645,27 +654,19 @@ async def google_assistant_update( msg: dict[str, Any], ) -> None: """Update google assistant entity config.""" - entity_registry = er.async_get(hass) entity_id: str = msg["entity_id"] - if not (registry_entry := entity_registry.async_get(entity_id)): - connection.send_error( - msg["id"], - websocket_api.const.ERR_NOT_ALLOWED, - f"can't configure {entity_id}", - ) - return + assistant_options: Mapping[str, Any] = {} + with suppress(HomeAssistantError, KeyError): + settings = exposed_entities.async_get_entity_settings(hass, entity_id) + assistant_options = settings[CLOUD_GOOGLE] disable_2fa = msg[PREF_DISABLE_2FA] - assistant_options: Mapping[str, Any] - if ( - assistant_options := registry_entry.options.get(CLOUD_GOOGLE, {}) - ) and assistant_options.get(PREF_DISABLE_2FA) == disable_2fa: + if assistant_options.get(PREF_DISABLE_2FA) == disable_2fa: return - assistant_options = assistant_options | {PREF_DISABLE_2FA: disable_2fa} - entity_registry.async_update_entity_options( - entity_id, CLOUD_GOOGLE, assistant_options + exposed_entities.async_set_assistant_option( + hass, CLOUD_GOOGLE, entity_id, PREF_DISABLE_2FA, disable_2fa ) connection.send_result(msg["id"]) diff --git a/homeassistant/components/homeassistant/exposed_entities.py b/homeassistant/components/homeassistant/exposed_entities.py index 159b54cb5a8..07f14e7ce8c 100644 --- a/homeassistant/components/homeassistant/exposed_entities.py +++ b/homeassistant/components/homeassistant/exposed_entities.py @@ -132,50 +132,52 @@ class ExposedEntities: self._listeners.setdefault(assistant, []).append(listener) @callback - def async_expose_entity( - self, assistant: str, entity_id: str, should_expose: bool + def async_set_assistant_option( + self, assistant: str, entity_id: str, key: str, value: Any ) -> None: - """Expose an entity to an assistant. + """Set an option for an assistant. Notify listeners if expose flag was changed. """ entity_registry = er.async_get(self._hass) if not (registry_entry := entity_registry.async_get(entity_id)): - return self._async_expose_legacy_entity(assistant, entity_id, should_expose) + return self._async_set_legacy_assistant_option( + assistant, entity_id, key, value + ) assistant_options: Mapping[str, Any] if ( assistant_options := registry_entry.options.get(assistant, {}) - ) and assistant_options.get("should_expose") == should_expose: + ) and assistant_options.get(key) == value: return - assistant_options = assistant_options | {"should_expose": should_expose} + assistant_options = assistant_options | {key: value} entity_registry.async_update_entity_options( entity_id, assistant, assistant_options ) for listener in self._listeners.get(assistant, []): listener() - def _async_expose_legacy_entity( - self, assistant: str, entity_id: str, should_expose: bool + def _async_set_legacy_assistant_option( + self, assistant: str, entity_id: str, key: str, value: Any ) -> None: - """Expose an entity to an assistant. + """Set an option for an assistant. Notify listeners if expose flag was changed. """ if ( (exposed_entity := self.entities.get(entity_id)) and (assistant_options := exposed_entity.assistants.get(assistant, {})) - and assistant_options.get("should_expose") == should_expose + and assistant_options.get(key) == value ): return if exposed_entity: new_exposed_entity = self._update_exposed_entity( - assistant, entity_id, should_expose + assistant, entity_id, key, value ) else: - new_exposed_entity = self._new_exposed_entity(assistant, should_expose) + new_exposed_entity = self._new_exposed_entity(assistant, key, value) self.entities[entity_id] = new_exposed_entity self._async_schedule_save() for listener in self._listeners.get(assistant, []): @@ -282,10 +284,12 @@ class ExposedEntities: if exposed_entity: new_exposed_entity = self._update_exposed_entity( - assistant, entity_id, should_expose + assistant, entity_id, "should_expose", should_expose ) else: - new_exposed_entity = self._new_exposed_entity(assistant, should_expose) + new_exposed_entity = self._new_exposed_entity( + assistant, "should_expose", should_expose + ) self.entities[entity_id] = new_exposed_entity self._async_schedule_save() @@ -322,22 +326,21 @@ class ExposedEntities: return False def _update_exposed_entity( - self, - assistant: str, - entity_id: str, - should_expose: bool, + self, assistant: str, entity_id: str, key: str, value: Any ) -> ExposedEntity: """Update an exposed entity.""" entity = self.entities[entity_id] assistants = dict(entity.assistants) old_settings = assistants.get(assistant, {}) - assistants[assistant] = old_settings | {"should_expose": should_expose} + assistants[assistant] = old_settings | {key: value} return ExposedEntity(assistants) - def _new_exposed_entity(self, assistant: str, should_expose: bool) -> ExposedEntity: + def _new_exposed_entity( + self, assistant: str, key: str, value: Any + ) -> ExposedEntity: """Create a new exposed entity.""" return ExposedEntity( - assistants={assistant: {"should_expose": should_expose}}, + assistants={assistant: {key: value}}, ) async def _async_load_data(self) -> SerializedExposedEntities | None: @@ -409,12 +412,9 @@ def ws_expose_entity( ) return - exposed_entities: ExposedEntities = hass.data[DATA_EXPOSED_ENTITIES] for entity_id in entity_ids: for assistant in msg["assistants"]: - exposed_entities.async_expose_entity( - assistant, entity_id, msg["should_expose"] - ) + async_expose_entity(hass, assistant, entity_id, msg["should_expose"]) connection.send_result(msg["id"]) @@ -513,8 +513,9 @@ def async_expose_entity( should_expose: bool, ) -> None: """Get assistant expose settings for an entity.""" - exposed_entities: ExposedEntities = hass.data[DATA_EXPOSED_ENTITIES] - exposed_entities.async_expose_entity(assistant, entity_id, should_expose) + async_set_assistant_option( + hass, assistant, entity_id, "should_expose", should_expose + ) @callback @@ -522,3 +523,15 @@ def async_should_expose(hass: HomeAssistant, assistant: str, entity_id: str) -> """Return True if an entity should be exposed to an assistant.""" exposed_entities: ExposedEntities = hass.data[DATA_EXPOSED_ENTITIES] return exposed_entities.async_should_expose(assistant, entity_id) + + +@callback +def async_set_assistant_option( + hass: HomeAssistant, assistant: str, entity_id: str, option: str, value: Any +) -> None: + """Set an option for an assistant. + + Notify listeners if expose flag was changed. + """ + exposed_entities: ExposedEntities = hass.data[DATA_EXPOSED_ENTITIES] + exposed_entities.async_set_assistant_option(assistant, entity_id, option, value) diff --git a/tests/components/cloud/test_alexa_config.py b/tests/components/cloud/test_alexa_config.py index 257d04cc697..134838ff1ce 100644 --- a/tests/components/cloud/test_alexa_config.py +++ b/tests/components/cloud/test_alexa_config.py @@ -15,6 +15,7 @@ from homeassistant.components.cloud.prefs import CloudPreferences from homeassistant.components.homeassistant.exposed_entities import ( DATA_EXPOSED_ENTITIES, ExposedEntities, + async_expose_entity, ) from homeassistant.const import EntityCategory from homeassistant.core import HomeAssistant @@ -40,8 +41,7 @@ def expose_new(hass, expose_new): def expose_entity(hass, entity_id, should_expose): """Expose an entity to Alexa.""" - exposed_entities: ExposedEntities = hass.data[DATA_EXPOSED_ENTITIES] - exposed_entities.async_expose_entity("cloud.alexa", entity_id, should_expose) + async_expose_entity(hass, "cloud.alexa", entity_id, should_expose) async def test_alexa_config_expose_entity_prefs( diff --git a/tests/components/cloud/test_client.py b/tests/components/cloud/test_client.py index d1e1a8ce112..534456896b4 100644 --- a/tests/components/cloud/test_client.py +++ b/tests/components/cloud/test_client.py @@ -16,6 +16,7 @@ from homeassistant.components.cloud.const import ( from homeassistant.components.homeassistant.exposed_entities import ( DATA_EXPOSED_ENTITIES, ExposedEntities, + async_expose_entity, ) from homeassistant.const import CONTENT_TYPE_JSON from homeassistant.core import HomeAssistant, State @@ -267,9 +268,7 @@ async def test_google_config_expose_entity( assert gconf.should_expose(state) - exposed_entities.async_expose_entity( - "cloud.google_assistant", entity_entry.entity_id, False - ) + async_expose_entity(hass, "cloud.google_assistant", entity_entry.entity_id, False) assert not gconf.should_expose(state) diff --git a/tests/components/cloud/test_google_config.py b/tests/components/cloud/test_google_config.py index 877a6efaf05..c6e08f9e152 100644 --- a/tests/components/cloud/test_google_config.py +++ b/tests/components/cloud/test_google_config.py @@ -18,6 +18,7 @@ from homeassistant.components.google_assistant import helpers as ga_helpers from homeassistant.components.homeassistant.exposed_entities import ( DATA_EXPOSED_ENTITIES, ExposedEntities, + async_expose_entity, ) from homeassistant.const import EVENT_HOMEASSISTANT_STARTED, EntityCategory from homeassistant.core import CoreState, HomeAssistant, State @@ -48,10 +49,7 @@ def expose_new(hass, expose_new): def expose_entity(hass, entity_id, should_expose): """Expose an entity to Google.""" - exposed_entities: ExposedEntities = hass.data[DATA_EXPOSED_ENTITIES] - exposed_entities.async_expose_entity( - "cloud.google_assistant", entity_id, should_expose - ) + async_expose_entity(hass, "cloud.google_assistant", entity_id, should_expose) async def test_google_update_report_state( diff --git a/tests/components/cloud/test_http_api.py b/tests/components/cloud/test_http_api.py index ff4b9be4d3f..f497c2c108e 100644 --- a/tests/components/cloud/test_http_api.py +++ b/tests/components/cloud/test_http_api.py @@ -14,6 +14,7 @@ from homeassistant.components.alexa import errors as alexa_errors from homeassistant.components.alexa.entities import LightCapabilities from homeassistant.components.cloud.const import DOMAIN from homeassistant.components.google_assistant.helpers import GoogleEntity +from homeassistant.components.homeassistant import exposed_entities from homeassistant.core import HomeAssistant, State from homeassistant.helpers import entity_registry as er from homeassistant.setup import async_setup_component @@ -842,6 +843,7 @@ async def test_get_google_entity( response = await client.receive_json() assert response["success"] assert response["result"] == { + "disable_2fa": None, "entity_id": "light.kitchen", "might_2fa": False, "traits": ["action.devices.traits.OnOff"], @@ -853,6 +855,30 @@ async def test_get_google_entity( response = await client.receive_json() assert response["success"] assert response["result"] == { + "disable_2fa": None, + "entity_id": "cover.garage", + "might_2fa": True, + "traits": ["action.devices.traits.OpenClose"], + } + + # Set the disable 2fa flag + await client.send_json_auto_id( + { + "type": "cloud/google_assistant/entities/update", + "entity_id": "cover.garage", + "disable_2fa": True, + } + ) + response = await client.receive_json() + assert response["success"] + + await client.send_json_auto_id( + {"type": "cloud/google_assistant/entities/get", "entity_id": "cover.garage"} + ) + response = await client.receive_json() + assert response["success"] + assert response["result"] == { + "disable_2fa": True, "entity_id": "cover.garage", "might_2fa": True, "traits": ["action.devices.traits.OpenClose"], @@ -867,9 +893,6 @@ async def test_update_google_entity( mock_cloud_login, ) -> None: """Test that we can update config of a Google entity.""" - entry = entity_registry.async_get_or_create( - "light", "test", "unique", suggested_object_id="kitchen" - ) client = await hass_ws_client(hass) await client.send_json_auto_id( { @@ -885,16 +908,16 @@ async def test_update_google_entity( { "type": "homeassistant/expose_entity", "assistants": ["cloud.google_assistant"], - "entity_ids": [entry.entity_id], + "entity_ids": ["light.kitchen"], "should_expose": False, } ) response = await client.receive_json() assert response["success"] - assert entity_registry.async_get(entry.entity_id).options[ - "cloud.google_assistant" - ] == {"disable_2fa": False, "should_expose": False} + assert exposed_entities.async_get_entity_settings(hass, "light.kitchen") == { + "cloud.google_assistant": {"disable_2fa": False, "should_expose": False} + } async def test_list_alexa_entities( diff --git a/tests/components/conversation/__init__.py b/tests/components/conversation/__init__.py index 6eadb068054..df57c78c9aa 100644 --- a/tests/components/conversation/__init__.py +++ b/tests/components/conversation/__init__.py @@ -7,6 +7,7 @@ from homeassistant.components import conversation from homeassistant.components.homeassistant.exposed_entities import ( DATA_EXPOSED_ENTITIES, ExposedEntities, + async_expose_entity, ) from homeassistant.helpers import intent @@ -53,5 +54,4 @@ def expose_new(hass, expose_new): def expose_entity(hass, entity_id, should_expose): """Expose an entity to the default agent.""" - exposed_entities: ExposedEntities = hass.data[DATA_EXPOSED_ENTITIES] - exposed_entities.async_expose_entity(conversation.DOMAIN, entity_id, should_expose) + async_expose_entity(hass, conversation.DOMAIN, entity_id, should_expose) diff --git a/tests/components/google_assistant/test_smart_home.py b/tests/components/google_assistant/test_smart_home.py index 6455128fce8..f9ea356216f 100644 --- a/tests/components/google_assistant/test_smart_home.py +++ b/tests/components/google_assistant/test_smart_home.py @@ -451,8 +451,8 @@ async def test_execute( hass: HomeAssistant, report_state, on, brightness, value ) -> None: """Test an execute command.""" - await async_setup_component(hass, "light", {"light": {"platform": "demo"}}) await async_setup_component(hass, "homeassistant", {}) + await async_setup_component(hass, "light", {"light": {"platform": "demo"}}) await hass.async_block_till_done() await hass.services.async_call( diff --git a/tests/components/homeassistant/test_exposed_entities.py b/tests/components/homeassistant/test_exposed_entities.py index db82a696155..fd09bcee45a 100644 --- a/tests/components/homeassistant/test_exposed_entities.py +++ b/tests/components/homeassistant/test_exposed_entities.py @@ -8,6 +8,7 @@ from homeassistant.components.homeassistant.exposed_entities import ( ExposedEntity, async_expose_entity, async_get_assistant_settings, + async_get_entity_settings, async_listen_entity_updates, async_should_expose, ) @@ -101,10 +102,10 @@ async def test_load_preferences(hass: HomeAssistant) -> None: exposed_entities.async_set_expose_new_entities("test1", True) exposed_entities.async_set_expose_new_entities("test2", False) - exposed_entities.async_expose_entity("test1", "light.kitchen", True) - exposed_entities.async_expose_entity("test1", "light.living_room", True) - exposed_entities.async_expose_entity("test2", "light.kitchen", True) - exposed_entities.async_expose_entity("test2", "light.kitchen", True) + async_expose_entity(hass, "test1", "light.kitchen", True) + async_expose_entity(hass, "test1", "light.living_room", True) + async_expose_entity(hass, "test2", "light.kitchen", True) + async_expose_entity(hass, "test2", "light.kitchen", True) assert list(exposed_entities._assistants) == ["test1", "test2"] assert list(exposed_entities.entities) == ["light.kitchen", "light.living_room"] @@ -334,27 +335,24 @@ async def test_listen_updates( assert await async_setup_component(hass, "homeassistant", {}) await hass.async_block_till_done() - exposed_entities: ExposedEntities = hass.data[DATA_EXPOSED_ENTITIES] async_listen_entity_updates(hass, "cloud.alexa", listener) entry = entity_registry.async_get_or_create("climate", "test", "unique1") # Call for another assistant - listener not called - exposed_entities.async_expose_entity( - "cloud.google_assistant", entry.entity_id, True - ) + async_expose_entity(hass, "cloud.google_assistant", entry.entity_id, True) assert len(calls) == 0 # Call for our assistant - listener called - exposed_entities.async_expose_entity("cloud.alexa", entry.entity_id, True) + async_expose_entity(hass, "cloud.alexa", entry.entity_id, True) assert len(calls) == 1 # Settings not changed - listener not called - exposed_entities.async_expose_entity("cloud.alexa", entry.entity_id, True) + async_expose_entity(hass, "cloud.alexa", entry.entity_id, True) assert len(calls) == 1 # Settings changed - listener called - exposed_entities.async_expose_entity("cloud.alexa", entry.entity_id, False) + async_expose_entity(hass, "cloud.alexa", entry.entity_id, False) assert len(calls) == 2 @@ -367,19 +365,17 @@ async def test_get_assistant_settings( assert await async_setup_component(hass, "homeassistant", {}) await hass.async_block_till_done() - exposed_entities: ExposedEntities = hass.data[DATA_EXPOSED_ENTITIES] - entry = entity_registry.async_get_or_create("climate", "test", "unique1") assert async_get_assistant_settings(hass, "cloud.alexa") == {} - exposed_entities.async_expose_entity("cloud.alexa", entry.entity_id, True) - exposed_entities.async_expose_entity("cloud.alexa", "light.not_in_registry", True) + async_expose_entity(hass, "cloud.alexa", entry.entity_id, True) + async_expose_entity(hass, "cloud.alexa", "light.not_in_registry", True) assert async_get_assistant_settings(hass, "cloud.alexa") == snapshot assert async_get_assistant_settings(hass, "cloud.google_assistant") == snapshot with pytest.raises(HomeAssistantError): - exposed_entities.async_get_entity_settings("light.unknown") + async_get_entity_settings(hass, "light.unknown") @pytest.mark.parametrize(