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

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)
+        try:
-
+            settings = async_get_entity_settings(self.hass, state.entity_id)
-        registry_entry = entity_registry.async_get(state.entity_id)
+        except HomeAssistantError:
-        if not registry_entry:
             # 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):

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)):
+    assistant_options: Mapping[str, Any] = {}
-        connection.send_error(
+    with suppress(HomeAssistantError, KeyError):
-            msg["id"],
+        settings = exposed_entities.async_get_entity_settings(hass, entity_id)
-            websocket_api.const.ERR_NOT_ALLOWED,
+        assistant_options = settings[CLOUD_GOOGLE]
-            f"can't configure {entity_id}",
-        )
-        return
 
     disable_2fa = msg[PREF_DISABLE_2FA]
-    assistant_options: Mapping[str, Any]
+    if assistant_options.get(PREF_DISABLE_2FA) == disable_2fa:
-    if (
-        assistant_options := registry_entry.options.get(CLOUD_GOOGLE, {})
-    ) and assistant_options.get(PREF_DISABLE_2FA) == disable_2fa:
         return
 
-    assistant_options = assistant_options | {PREF_DISABLE_2FA: disable_2fa}
+    exposed_entities.async_set_assistant_option(
-    entity_registry.async_update_entity_options(
+        hass, CLOUD_GOOGLE, entity_id, PREF_DISABLE_2FA, disable_2fa
-        entity_id, CLOUD_GOOGLE, assistant_options
     )
     connection.send_result(msg["id"])
 

homeassistant/components/homeassistant/exposed_entities.py

@@ -132,50 +132,52 @@ class ExposedEntities:
         self._listeners.setdefault(assistant, []).append(listener)
 
     @callback
-    def async_expose_entity(
+    def async_set_assistant_option(
-        self, assistant: str, entity_id: str, should_expose: bool
+        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(
+    def _async_set_legacy_assistant_option(
-        self, assistant: str, entity_id: str, should_expose: bool
+        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,
+        self, assistant: str, entity_id: str, key: str, value: Any
-        assistant: str,
-        entity_id: str,
-        should_expose: bool,
     ) -> 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(
+            async_expose_entity(hass, assistant, entity_id, msg["should_expose"])
-                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]
+    async_set_assistant_option(
-    exposed_entities.async_expose_entity(assistant, entity_id, should_expose)
+        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)

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]
+    async_expose_entity(hass, "cloud.alexa", entity_id, should_expose)
-    exposed_entities.async_expose_entity("cloud.alexa", entity_id, should_expose)
 
 
 async def test_alexa_config_expose_entity_prefs(

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(
+    async_expose_entity(hass, "cloud.google_assistant", entity_entry.entity_id, False)
-        "cloud.google_assistant", entity_entry.entity_id, False
-    )
 
     assert not gconf.should_expose(state)
 

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]
+    async_expose_entity(hass, "cloud.google_assistant", entity_id, should_expose)
-    exposed_entities.async_expose_entity(
-        "cloud.google_assistant", entity_id, should_expose
-    )
 
 
 async def test_google_update_report_state(

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[
+    assert exposed_entities.async_get_entity_settings(hass, "light.kitchen") == {
-        "cloud.google_assistant"
+        "cloud.google_assistant": {"disable_2fa": False, "should_expose": False}
-    ] == {"disable_2fa": False, "should_expose": False}
+    }
 
 
 async def test_list_alexa_entities(

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]
+    async_expose_entity(hass, conversation.DOMAIN, entity_id, should_expose)
-    exposed_entities.async_expose_entity(conversation.DOMAIN, entity_id, should_expose)

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(

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)
+    async_expose_entity(hass, "test1", "light.kitchen", True)
-    exposed_entities.async_expose_entity("test1", "light.living_room", True)
+    async_expose_entity(hass, "test1", "light.living_room", True)
-    exposed_entities.async_expose_entity("test2", "light.kitchen", True)
+    async_expose_entity(hass, "test2", "light.kitchen", True)
-    exposed_entities.async_expose_entity("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(
+    async_expose_entity(hass, "cloud.google_assistant", entry.entity_id, True)
-        "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)
+    async_expose_entity(hass, "cloud.alexa", entry.entity_id, True)
-    exposed_entities.async_expose_entity("cloud.alexa", "light.not_in_registry", 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(