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
This commit is contained in:
Erik Montnemery 2023-05-03 15:55:38 +02:00 committed by GitHub
parent 20942ab26f
commit 31de1b17e8
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
10 changed files with 113 additions and 82 deletions

View File

@ -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):

View File

@ -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"])

View File

@ -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)

View File

@ -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(

View File

@ -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)

View File

@ -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(

View File

@ -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(

View File

@ -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)

View File

@ -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(

View File

@ -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(