Use HassKey in homeassistant integration (#120332)

homeassistant/components/homeassistant/const.py

@@ -1,12 +1,18 @@
 """Constants for the Homeassistant integration."""
 
-from typing import Final
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Final
 
 import homeassistant.core as ha
+from homeassistant.util.hass_dict import HassKey
+
+if TYPE_CHECKING:
+    from .exposed_entities import ExposedEntities
 
 DOMAIN = ha.DOMAIN
 
-DATA_EXPOSED_ENTITIES = f"{DOMAIN}.exposed_entites"
+DATA_EXPOSED_ENTITIES: HassKey[ExposedEntities] = HassKey(f"{DOMAIN}.exposed_entites")
 DATA_STOP_HANDLER = f"{DOMAIN}.stop_handler"
 
 SERVICE_HOMEASSISTANT_STOP: Final = "stop"

homeassistant/components/homeassistant/exposed_entities.py

@@ -440,7 +440,7 @@ def ws_list_exposed_entities(
     """Expose an entity to an assistant."""
     result: dict[str, Any] = {}
 
-    exposed_entities: ExposedEntities = hass.data[DATA_EXPOSED_ENTITIES]
+    exposed_entities = hass.data[DATA_EXPOSED_ENTITIES]
     entity_registry = er.async_get(hass)
     for entity_id in chain(exposed_entities.entities, entity_registry.entities):
         result[entity_id] = {}
@@ -464,7 +464,7 @@ def ws_expose_new_entities_get(
     hass: HomeAssistant, connection: websocket_api.ActiveConnection, msg: dict[str, Any]
 ) -> None:
     """Check if new entities are exposed to an assistant."""
-    exposed_entities: ExposedEntities = hass.data[DATA_EXPOSED_ENTITIES]
+    exposed_entities = hass.data[DATA_EXPOSED_ENTITIES]
     expose_new = exposed_entities.async_get_expose_new_entities(msg["assistant"])
     connection.send_result(msg["id"], {"expose_new": expose_new})
 
@@ -482,7 +482,7 @@ def ws_expose_new_entities_set(
     hass: HomeAssistant, connection: websocket_api.ActiveConnection, msg: dict[str, Any]
 ) -> None:
     """Expose new entities to an assistant."""
-    exposed_entities: ExposedEntities = hass.data[DATA_EXPOSED_ENTITIES]
+    exposed_entities = hass.data[DATA_EXPOSED_ENTITIES]
     exposed_entities.async_set_expose_new_entities(msg["assistant"], msg["expose_new"])
     connection.send_result(msg["id"])
 
@@ -492,7 +492,7 @@ def async_listen_entity_updates(
     hass: HomeAssistant, assistant: str, listener: Callable[[], None]
 ) -> CALLBACK_TYPE:
     """Listen for updates to entity expose settings."""
-    exposed_entities: ExposedEntities = hass.data[DATA_EXPOSED_ENTITIES]
+    exposed_entities = hass.data[DATA_EXPOSED_ENTITIES]
     return exposed_entities.async_listen_entity_updates(assistant, listener)
 
 
@@ -501,7 +501,7 @@ def async_get_assistant_settings(
     hass: HomeAssistant, assistant: str
 ) -> dict[str, Mapping[str, Any]]:
     """Get all entity expose settings for an assistant."""
-    exposed_entities: ExposedEntities = hass.data[DATA_EXPOSED_ENTITIES]
+    exposed_entities = hass.data[DATA_EXPOSED_ENTITIES]
     return exposed_entities.async_get_assistant_settings(assistant)
 
 
@@ -510,7 +510,7 @@ def async_get_entity_settings(
     hass: HomeAssistant, entity_id: str
 ) -> dict[str, Mapping[str, Any]]:
     """Get assistant expose settings for an entity."""
-    exposed_entities: ExposedEntities = hass.data[DATA_EXPOSED_ENTITIES]
+    exposed_entities = hass.data[DATA_EXPOSED_ENTITIES]
     return exposed_entities.async_get_entity_settings(entity_id)
 
 
@@ -530,7 +530,7 @@ def async_expose_entity(
 @callback
 def async_should_expose(hass: HomeAssistant, assistant: str, entity_id: str) -> bool:
     """Return True if an entity should be exposed to an assistant."""
-    exposed_entities: ExposedEntities = hass.data[DATA_EXPOSED_ENTITIES]
+    exposed_entities = hass.data[DATA_EXPOSED_ENTITIES]
     return exposed_entities.async_should_expose(assistant, entity_id)
 
 
@@ -542,5 +542,5 @@ def async_set_assistant_option(
 
     Notify listeners if expose flag was changed.
     """
-    exposed_entities: ExposedEntities = hass.data[DATA_EXPOSED_ENTITIES]
+    exposed_entities = hass.data[DATA_EXPOSED_ENTITIES]
     exposed_entities.async_set_assistant_option(assistant, entity_id, option, value)

tests/components/cloud/test_alexa_config.py

@@ -15,7 +15,6 @@ from homeassistant.components.cloud.const import (
 from homeassistant.components.cloud.prefs import CloudPreferences
 from homeassistant.components.homeassistant.exposed_entities import (
     DATA_EXPOSED_ENTITIES,
-    ExposedEntities,
     async_expose_entity,
     async_get_entity_settings,
 )
@@ -39,13 +38,13 @@ def cloud_stub():
     return Mock(is_logged_in=True, subscription_expired=False)
 
 
-def expose_new(hass, expose_new):
+def expose_new(hass: HomeAssistant, expose_new: bool) -> None:
     """Enable exposing new entities to Alexa."""
-    exposed_entities: ExposedEntities = hass.data[DATA_EXPOSED_ENTITIES]
+    exposed_entities = hass.data[DATA_EXPOSED_ENTITIES]
     exposed_entities.async_set_expose_new_entities("cloud.alexa", expose_new)
 
 
-def expose_entity(hass, entity_id, should_expose):
+def expose_entity(hass: HomeAssistant, entity_id: str, should_expose: bool) -> None:
     """Expose an entity to Alexa."""
     async_expose_entity(hass, "cloud.alexa", entity_id, should_expose)
 

tests/components/cloud/test_client.py

@@ -21,7 +21,6 @@ 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, __version__ as HA_VERSION
@@ -262,7 +261,7 @@ async def test_google_config_expose_entity(
     """Test Google config exposing entity method uses latest config."""
 
     # Enable exposing new entities to Google
-    exposed_entities: ExposedEntities = hass.data[DATA_EXPOSED_ENTITIES]
+    exposed_entities = hass.data[DATA_EXPOSED_ENTITIES]
     exposed_entities.async_set_expose_new_entities("cloud.google_assistant", True)
 
     # Register a light entity

tests/components/cloud/test_google_config.py

@@ -18,7 +18,6 @@ from homeassistant.components.cloud.prefs import CloudPreferences
 from homeassistant.components.google_assistant import helpers as ga_helpers
 from homeassistant.components.homeassistant.exposed_entities import (
     DATA_EXPOSED_ENTITIES,
-    ExposedEntities,
     async_expose_entity,
     async_get_entity_settings,
 )
@@ -47,13 +46,13 @@ def mock_conf(hass, cloud_prefs):
     )
 
 
-def expose_new(hass, expose_new):
+def expose_new(hass: HomeAssistant, expose_new: bool) -> None:
     """Enable exposing new entities to Google."""
-    exposed_entities: ExposedEntities = hass.data[DATA_EXPOSED_ENTITIES]
+    exposed_entities = hass.data[DATA_EXPOSED_ENTITIES]
     exposed_entities.async_set_expose_new_entities("cloud.google_assistant", expose_new)
 
 
-def expose_entity(hass, entity_id, should_expose):
+def expose_entity(hass: HomeAssistant, entity_id: str, should_expose: bool) -> None:
     """Expose an entity to Google."""
     async_expose_entity(hass, "cloud.google_assistant", entity_id, should_expose)
 

tests/components/conversation/__init__.py

@@ -11,7 +11,6 @@ from homeassistant.components.conversation.models import (
 )
 from homeassistant.components.homeassistant.exposed_entities import (
     DATA_EXPOSED_ENTITIES,
-    ExposedEntities,
     async_expose_entity,
 )
 from homeassistant.core import HomeAssistant
@@ -45,12 +44,12 @@ class MockAgent(conversation.AbstractConversationAgent):
         )
 
 
-def expose_new(hass: HomeAssistant, expose_new: bool):
+def expose_new(hass: HomeAssistant, expose_new: bool) -> None:
     """Enable exposing new entities to the default agent."""
-    exposed_entities: ExposedEntities = hass.data[DATA_EXPOSED_ENTITIES]
+    exposed_entities = hass.data[DATA_EXPOSED_ENTITIES]
     exposed_entities.async_set_expose_new_entities(conversation.DOMAIN, expose_new)
 
 
-def expose_entity(hass: HomeAssistant, entity_id: str, should_expose: bool):
+def expose_entity(hass: HomeAssistant, entity_id: str, should_expose: bool) -> None:
     """Expose an entity to the default agent."""
     async_expose_entity(hass, conversation.DOMAIN, entity_id, should_expose)

tests/components/homeassistant/test_exposed_entities.py

@@ -103,7 +103,7 @@ async def test_load_preferences(hass: HomeAssistant) -> None:
     """Make sure that we can load/save data correctly."""
     assert await async_setup_component(hass, "homeassistant", {})
 
-    exposed_entities: ExposedEntities = hass.data[DATA_EXPOSED_ENTITIES]
+    exposed_entities = hass.data[DATA_EXPOSED_ENTITIES]
     assert exposed_entities._assistants == {}
 
     exposed_entities.async_set_expose_new_entities("test1", True)
@@ -139,7 +139,7 @@ async def test_expose_entity(
     entry1 = entity_registry.async_get_or_create("test", "test", "unique1")
     entry2 = entity_registry.async_get_or_create("test", "test", "unique2")
 
-    exposed_entities: ExposedEntities = hass.data[DATA_EXPOSED_ENTITIES]
+    exposed_entities = hass.data[DATA_EXPOSED_ENTITIES]
     assert len(exposed_entities.entities) == 0
 
     # Set options
@@ -196,7 +196,7 @@ async def test_expose_entity_unknown(
     assert await async_setup_component(hass, "homeassistant", {})
     await hass.async_block_till_done()
 
-    exposed_entities: ExposedEntities = hass.data[DATA_EXPOSED_ENTITIES]
+    exposed_entities = hass.data[DATA_EXPOSED_ENTITIES]
     assert len(exposed_entities.entities) == 0
 
     # Set options
@@ -442,7 +442,7 @@ async def test_should_expose(
     )
 
     # Check with a different assistant
-    exposed_entities: ExposedEntities = hass.data[DATA_EXPOSED_ENTITIES]
+    exposed_entities = hass.data[DATA_EXPOSED_ENTITIES]
     exposed_entities.async_set_expose_new_entities("cloud.no_default_expose", False)
     assert (
         async_should_expose(
@@ -545,7 +545,7 @@ async def test_listeners(
     """Make sure we call entity listeners."""
     assert await async_setup_component(hass, "homeassistant", {})
 
-    exposed_entities: ExposedEntities = hass.data[DATA_EXPOSED_ENTITIES]
+    exposed_entities = hass.data[DATA_EXPOSED_ENTITIES]
 
     callbacks = []
     exposed_entities.async_listen_entity_updates("test1", lambda: callbacks.append(1))