Merge branch 'dev' into tests-sensor-tracker-text

Moved to its own pull request.

esphome/components/as734x/sensor.py

@@ -108,7 +108,8 @@ CONFIG_SCHEMA = cv.Schema(
 
 
 def _final_validate(config: ConfigType) -> ConfigType:
-    hub = fv.get_declaring_config(config[CONF_AS734X_ID])
+    full = fv.full_config.get()
+    hub = full.get_config_for_path(full.get_path_for_id(config[CONF_AS734X_ID])[:-1])
     if hub[CONF_TYPE] != "AS7341":
         return config
     for key in _BAND_KEYS:

esphome/components/bmi270/sensor.py

@@ -66,7 +66,9 @@ CONFIG_SCHEMA = cv.typed_schema(
 def _final_validate(config: ConfigType) -> None:
     if config[CONF_TYPE] not in _MAGNETIC_FIELDS:
         return
-    bmi270_config = fv.get_declaring_config(config[CONF_BMI270_ID])
+    full_config = fv.full_config.get()
+    bmi270_path = full_config.get_path_for_id(config[CONF_BMI270_ID])[:-1]
+    bmi270_config = full_config.get_config_for_path(bmi270_path)
     if bmi270_config.get(CONF_AUX_DEVICE) != "BMM150":
         raise cv.Invalid(
             "This sensor requires the parent bmi270's 'aux_device' to be set to 'BMM150'",

esphome/components/bq25186/__init__.py

@@ -7,7 +7,6 @@ from esphome.components import i2c
 import esphome.config_validation as cv
 from esphome.const import CONF_ID, CONF_UPDATE_INTERVAL, CONF_WATCHDOG
 from esphome.core import CORE, ID, EsphomeError, Lambda, TimePeriod
-import esphome.final_validate as fv
 from esphome.types import ConfigType, TemplateArgsType
 
 CODEOWNERS = ["@Rapsssito"]
@@ -360,7 +359,8 @@ async def bq25186_configure_to_code(
     parent = await cg.get_variable(config[CONF_ID])
 
     # Options the action leaves out fall back to the component's configuration
-    parent_config = fv.get_declaring_config(config[CONF_ID], CORE.config)
+    path = CORE.config.get_path_for_id(config[CONF_ID])[:-1]
+    parent_config = CORE.config.get_config_for_path(path)
     merged = {**parent_config, **config}
     try:
         _check_watchdog(merged)

esphome/components/dfrobot_sen0609/__init__.py

@@ -3,7 +3,8 @@ from collections.abc import Callable
 import esphome.codegen as cg
 from esphome.components import uart
 import esphome.config_validation as cv
-from esphome.const import CONF_MODE, CONF_MODEL
+from esphome.const import CONF_ID, CONF_MODE, CONF_MODEL
+from esphome.core import ID
 import esphome.final_validate as fv
 from esphome.types import ConfigType
 
@@ -58,11 +59,16 @@ FINAL_VALIDATE_SCHEMA = uart.final_validate_device_schema(
 )
 
 
+def get_hub_config(root: ConfigType, hub_id: ID) -> ConfigType:
+    """Return the configuration of the hub with the given id."""
+    return next(conf for conf in root[DOMAIN] if conf[CONF_ID] == hub_id)
+
+
 def final_validate_mode(keys: dict[str, str]) -> Callable[[ConfigType], None]:
     """Reject entity keys that need a different hub mode; `keys` maps each key to the mode it needs."""
 
     def validator(config: ConfigType) -> None:
-        hub = fv.get_declaring_config(config[CONF_DFROBOT_SEN0609_ID])
+        hub = get_hub_config(fv.full_config.get(), config[CONF_DFROBOT_SEN0609_ID])
         for key, mode in keys.items():
             if key in config and hub[CONF_MODE] != mode:
                 raise cv.Invalid(

esphome/components/dfrobot_sen0609/number/__init__.py

@@ -18,7 +18,6 @@ from esphome.const import (
     UNIT_SECOND,
 )
 from esphome.core import CORE
-import esphome.final_validate as fv
 from esphome.types import ConfigType
 
 from .. import (
@@ -29,6 +28,7 @@ from .. import (
     DFRobotSEN0609Component,
     dfrobot_sen0609_ns,
     final_validate_mode,
+    get_hub_config,
 )
 
 DFRobotSEN0609Number = dfrobot_sen0609_ns.class_("DFRobotSEN0609Number", number.Number)
@@ -181,7 +181,7 @@ FINAL_VALIDATE_SCHEMA = final_validate_mode(
 async def to_code(config: ConfigType) -> None:
     hub_id = config[CONF_DFROBOT_SEN0609_ID]
     hub = await cg.get_variable(hub_id)
-    max_range = MAX_RANGES[fv.get_declaring_config(hub_id, CORE.config)[CONF_MODEL]]
+    max_range = MAX_RANGES[get_hub_config(CORE.config, hub_id)[CONF_MODEL]]
     numbers = number.sub_numbers(config, parent=hub)
     for s in SETTINGS:
         await numbers(

esphome/components/display/__init__.py

@@ -26,7 +26,7 @@ from esphome.const import (
     SCHEDULER_DONT_RUN,
 )
 from esphome.core import CORE, ID, CoroPriority, coroutine_with_priority
-from esphome.final_validate import full_config, get_declaring_config
+from esphome.final_validate import full_config
 from esphome.types import ConfigType
 
 DOMAIN = "display"
@@ -196,7 +196,9 @@ def get_display_metadata(display_id: ID) -> DisplayMetaData:
             return meta
     # No metadata found, display driver may not yet support it.
     # Read the raw config to populate the returned data
-    disp_config = get_declaring_config(display_id)
+    global_config = full_config.get()
+    path = global_config.get_path_for_id(display_id)[:-1]
+    disp_config = global_config.get_config_for_path(path)
     dimensions = disp_config.get(CONF_DIMENSIONS, (0, 0))
     if isinstance(dimensions, dict):
         dimensions = (dimensions.get(CONF_WIDTH, 0), dimensions.get(CONF_HEIGHT, 0))

esphome/components/ds248x/one_wire.py

@@ -32,7 +32,9 @@ CONFIG_SCHEMA = cv.Schema(
 
 def _final_validate(config: ConfigType) -> None:
     """Validate that the channel is within the parent's channel count."""
-    parent_config = fv.get_declaring_config(config[CONF_DS248X_ID])
+    fconf = fv.full_config.get()
+    path = fconf.get_path_for_id(config[CONF_DS248X_ID])[:-1]
+    parent_config = fconf.get_config_for_path(path)
     channel_count = get_channel_count(parent_config)
     channel = config[CONF_CHANNEL]
 

esphome/components/gree/switch/__init__.py

@@ -63,7 +63,9 @@ CONFIG_SCHEMA = cv.Schema(
 
 
 def _validate_model(config: ConfigType) -> None:
-    climate_conf = fv.get_declaring_config(config[CONF_GREE_ID])
+    full_config = fv.full_config.get()
+    climate_path = full_config.get_path_for_id(config[CONF_GREE_ID])[:-1]
+    climate_conf = full_config.get_config_for_path(climate_path)
     if climate_conf[CONF_MODEL] not in SUPPORTED_MODELS:
         raise cv.Invalid(
             "Gree switches are only supported for the "

esphome/components/haier/switch/__init__.py

@@ -61,10 +61,12 @@ CONFIG_SCHEMA = cv.Schema(
 
 
 def _final_validate(config) -> None:
+    full_config = fv.full_config.get()
     for switch_type in [CONF_BEEPER, CONF_QUIET_MODE]:
         # Check switches that are only supported for HonClimate
         if config.get(switch_type):
-            climate_conf = fv.get_declaring_config(config[CONF_HAIER_ID])
+            climate_path = full_config.get_path_for_id(config[CONF_HAIER_ID])[:-1]
+            climate_conf = full_config.get_config_for_path(climate_path)
             protocol_type = climate_conf.get(CONF_PROTOCOL)
             if protocol_type.casefold() != PROTOCOL_HON.casefold():
                 raise cv.Invalid(

esphome/components/ir_rf_proxy/infrared.py

@@ -42,7 +42,9 @@ def _final_validate(config: dict[str, Any]) -> None:
 
     # Get the transmitter configuration
     transmitter_id = config[CONF_REMOTE_TRANSMITTER_ID]
-    transmitter_config = fv.get_declaring_config(transmitter_id)
+    full_config = fv.full_config.get()
+    transmitter_path = full_config.get_path_for_id(transmitter_id)[:-1]
+    transmitter_config = full_config.get_config_for_path(transmitter_path)
 
     # Check if carrier_duty_percent set to 0 or 100
     # Note: remote_transmitter schema requires this field and validates 1-100%,

esphome/components/ir_rf_proxy/radio_frequency.py

@@ -35,7 +35,11 @@ def _final_validate(config: ConfigType) -> None:
     if CONF_REMOTE_TRANSMITTER_ID not in config:
         return
 
-    transmitter_config = fv.get_declaring_config(config[CONF_REMOTE_TRANSMITTER_ID])
+    full_config = fv.full_config.get()
+    transmitter_path = full_config.get_path_for_id(config[CONF_REMOTE_TRANSMITTER_ID])[
+        :-1
+    ]
+    transmitter_config = full_config.get_config_for_path(transmitter_path)
 
     duty_percent = transmitter_config.get(CONF_CARRIER_DUTY_PERCENT)
     if duty_percent is not None and duty_percent != 100:

esphome/components/ld6002b/button/__init__.py

@@ -101,7 +101,8 @@ def final_validate(config: ConfigType) -> None:
             )
 
     if config.get(CONF_WAKE):
-        hub_config = fv.get_declaring_config(hub_id, full_config)
+        hub_path = full_config.get_path_for_id(hub_id)
+        hub_config = full_config.get_config_for_path(hub_path[:-1])
         if hub_config.get(CONF_WAKEUP_PIN) is None:
             raise cv.Invalid(
                 f"{CONF_WAKE} requires {CONF_WAKEUP_PIN} on the parent ld6002b component",

esphome/components/light/__init__.py

@@ -380,7 +380,8 @@ def _final_validate(config: ConfigType) -> None:
 
     for ref in refs:
         try:
-            light_config = fv.get_declaring_config(ref.light_id, fconf)
+            light_path = fconf.get_path_for_id(ref.light_id)[:-1]
+            light_config = fconf.get_config_for_path(light_path)
         except KeyError:
             # Light ID not found — ID validation will have already reported this
             continue
@@ -397,7 +398,8 @@ def _final_validate(config: ConfigType) -> None:
 
     for ref in cycle_refs:
         try:
-            light_config = fv.get_declaring_config(ref.light_id, fconf)
+            light_path = fconf.get_path_for_id(ref.light_id)[:-1]
+            light_config = fconf.get_config_for_path(light_path)
         except KeyError:
             continue
 

esphome/components/light/automation.py

@@ -34,7 +34,6 @@ from esphome.const import (
 )
 from esphome.core import CORE, ID, EsphomeError, Lambda
 from esphome.cpp_generator import MockObj, TemplateArgsType
-import esphome.final_validate as fv
 from esphome.schema_extractors import SCHEMA_EXTRACT, schema_extractor
 from esphome.types import ConfigType
 
@@ -254,7 +253,8 @@ def _resolve_effect_index(config: ConfigType, original_name: str) -> int:
     from . import effect_not_found_message, find_effect_index
 
     light_id = config[CONF_ID]
-    light_config = fv.get_declaring_config(light_id, CORE.config)
+    light_path = CORE.config.get_path_for_id(light_id)[:-1]
+    light_config = CORE.config.get_config_for_path(light_path)
     effects = light_config.get(CONF_EFFECTS, [])
     index = find_effect_index(effects, original_name)
     if index is not None:

esphome/components/mipi_spi/light/__init__.py

@@ -41,7 +41,9 @@ CONFIG_SCHEMA = cv.All(
 
 
 def _final_validate(config: ConfigType) -> None:
-    display_config = fv.get_declaring_config(config[CONF_DISPLAY_ID])
+    full_config = fv.full_config.get()
+    display_path = full_config.get_path_for_id(config[CONF_DISPLAY_ID])[:-1]
+    display_config = full_config.get_config_for_path(display_path)
     if CONF_BRIGHTNESS not in display_config:
         model = display_config[CONF_MODEL]
         if model == "CUSTOM":

esphome/components/modbus_controller/__init__.py

@@ -310,7 +310,9 @@ def migrate_custom_command(config: ConfigType) -> None:
     frame = config.get(CONF_CUSTOM_COMMAND)
     if frame is None:
         return
-    controller = fv.get_declaring_config(config[CONF_MODBUS_CONTROLLER_ID])
+    fconf = fv.full_config.get()
+    path = fconf.get_path_for_id(config[CONF_MODBUS_CONTROLLER_ID])[:-1]
+    controller = fconf.get_config_for_path(path)
     # the controller's DEVICE address (from modbus_device_schema)
     address = controller[CONF_ADDRESS]
     if frame[0] != address:
@@ -340,7 +342,9 @@ def _reject_continuous_write_custom_pdu(config: ConfigType) -> None:
     pdu = config.get(CONF_CUSTOM_PDU)
     if pdu is None or not modbus.is_function_code_write(pdu[0]):
         return
-    controller = fv.get_declaring_config(config[CONF_MODBUS_CONTROLLER_ID])
+    fconf = fv.full_config.get()
+    path = fconf.get_path_for_id(config[CONF_MODBUS_CONTROLLER_ID])[:-1]
+    controller = fconf.get_config_for_path(path)
     if controller.get(CONF_CONTINUOUS) is True:
         raise cv.Invalid(
             f"a '{CONF_CUSTOM_PDU}' with a write function code (0x{pdu[0] & 0x7F:02X}) can't be polled "
@@ -355,7 +359,9 @@ def _reject_broadcastable_custom_pdu(config: ConfigType) -> None:
     pdu = config.get(CONF_CUSTOM_PDU)
     if pdu is None or not modbus.is_function_code_broadcastable(pdu[0]):
         return
-    controller = fv.get_declaring_config(config[CONF_MODBUS_CONTROLLER_ID])
+    fconf = fv.full_config.get()
+    path = fconf.get_path_for_id(config[CONF_MODBUS_CONTROLLER_ID])[:-1]
+    controller = fconf.get_config_for_path(path)
     if (
         controller.get(CONF_ADDRESS) == modbus.BROADCAST_ADDRESS
         and controller.get(modbus.CONF_ALLOW_BROADCAST_READ) is True
@@ -378,7 +384,9 @@ def validate_custom_pdu_item(config: ConfigType) -> None:
 def _reject_write_option_off_broadcast(config: ConfigType) -> None:
     if not any(config.get(key) is True for key in modbus.broadcast_only_option_keys()):
         return
-    controller = fv.get_declaring_config(config[CONF_MODBUS_CONTROLLER_ID])
+    fconf = fv.full_config.get()
+    path = fconf.get_path_for_id(config[CONF_MODBUS_CONTROLLER_ID])[:-1]
+    controller = fconf.get_config_for_path(path)
     if controller.get(CONF_ADDRESS) != modbus.BROADCAST_ADDRESS:
         raise cv.Invalid(
             f"'{modbus.CONF_EXPECT_BROADCAST_WRITE_RESPONSE}' only applies when the "

esphome/components/motion/__init__.py

@@ -50,7 +50,9 @@ MAX_RECOMMENDED_UPDATE_INTERVAL_MS = 100
 
 def get_motion_config(motion_id: ID) -> dict:
     """Look up the (already-validated) config of the motion hub referenced by motion_id."""
-    return fv.get_declaring_config(motion_id)
+    full_config = fv.full_config.get()
+    motion_path = full_config.get_path_for_id(motion_id)[:-1]
+    return full_config.get_config_for_path(motion_path)
 
 
 def check_update_interval(motion_id: ID, feature_name: str) -> None:

esphome/components/mpr121/__init__.py

@@ -67,7 +67,8 @@ def _final_validate(config: ConfigType) -> None:
             "Max touch channel must equal the highest binary sensor channel or be removed for auto calculation",
             path=[CONF_MAX_TOUCH_CHANNEL],
         )
-    this_config = fv.get_declaring_config(config[CONF_ID], fconf)
+    path = fconf.get_path_for_id(config[CONF_ID])[:-1]
+    this_config = fconf.get_config_for_path(path)
     this_config[CONF_MAX_TOUCH_CHANNEL] = max_touch_channel
 
 

esphome/components/packet_transport/binary_sensor.py

@@ -48,7 +48,9 @@ def _final_validate(config: ConfigType) -> None:
     if config[CONF_TYPE] != CONF_STATUS:
         # Only run this validation if a status sensor is being configured
         return
-    transport_config = fv.get_declaring_config(config[CONF_TRANSPORT_ID])
+    full_config = fv.full_config.get()
+    transport_path = full_config.get_path_for_id(config[CONF_TRANSPORT_ID])[:-1]
+    transport_config = full_config.get_config_for_path(transport_path)
     if transport_config[CONF_PING_PONG_ENABLE] and any(
         CONF_ENCRYPTION in p
         for p in transport_config[CONF_PROVIDERS]

esphome/components/partition/light.py

@@ -35,8 +35,11 @@ def validate_from_to(value: ConfigType) -> ConfigType:
 
 
 def validate_segment(config: ConfigType) -> None:
+    fconf = fv.full_config.get()
+
     if CONF_ID in config:  # only validate addressable segments
-        segment_light_config = fv.get_declaring_config(config[CONF_ID])
+        path = fconf.get_path_for_id(config[CONF_ID])[:-1]
+        segment_light_config = fconf.get_config_for_path(path)
 
         if CONF_NUM_LEDS in segment_light_config:
             segment_len = segment_light_config[CONF_NUM_LEDS]

esphome/components/tas58xx/binary_sensor.py

@@ -56,7 +56,9 @@ CONFIG_SCHEMA = cv.Schema(
 
 
 def _final_validate(config: ConfigType) -> ConfigType:
-    hub_conf = fv.get_declaring_config(config[CONF_TAS58XX_ID])
+    fconf = fv.full_config.get()
+    hub_path = fconf.get_path_for_id(config[CONF_TAS58XX_ID])
+    hub_conf = fconf.get_config_for_path(hub_path[:-1])
 
     if hub_conf[CONF_MODEL] == DAC_TAS5825M:
         return config

esphome/components/tca8418/binary_sensor/__init__.py

@@ -65,7 +65,9 @@ CONFIG_SCHEMA = cv.All(
 
 def _final_validate(config: ConfigType) -> None:
     """Resolve the key to the number the device reports, checked against the keypad's matrix."""
-    keypad = fv.get_declaring_config(config[CONF_KEYPAD_ID])
+    full_config = fv.full_config.get()
+    keypad_path = full_config.get_path_for_id(config[CONF_KEYPAD_ID])[:-1]
+    keypad = full_config.get_config_for_path(keypad_path)
     rows = keypad.get(CONF_ROWS, 0)
     columns = keypad.get(CONF_COLUMNS, 0)
 

esphome/components/tmc2209/sensor/__init__.py

@@ -59,7 +59,9 @@ CONFIG_SCHEMA = cv.Schema(
 def _final_validate(config: ConfigType) -> ConfigType:
     if CONF_ACTUAL_CURRENT not in config:
         return config
-    if fv.get_declaring_config(config[CONF_TMC2209_ID]).get(CONF_ANALOG_CURRENT_SCALE):
+    full_config = fv.full_config.get()
+    path = full_config.get_path_for_id(config[CONF_TMC2209_ID])[:-1]
+    if full_config.get_config_for_path(path).get(CONF_ANALOG_CURRENT_SCALE):
         raise cv.Invalid(
             f"'{CONF_ACTUAL_CURRENT}' can't be used with '{CONF_ANALOG_CURRENT_SCALE}', "
             "the current then depends on VREF",

esphome/components/tmp102/number.py

@@ -56,7 +56,8 @@ _validate_unique = unique_child("number")
 
 def _final_validate(config: ConfigType) -> ConfigType:
     _validate_unique(config)
-    parent = fv.get_declaring_config(config[CONF_TMP102_ID])
+    full = fv.full_config.get()
+    parent = full.get_config_for_path(full.get_path_for_id(config[CONF_TMP102_ID])[:-1])
     add_number_parent(config[CONF_TMP102_ID], parent.get(CONF_EXTENDED_MODE, False))
     return config
 

esphome/config.py

@@ -1136,7 +1136,9 @@ class IDPassValidationStep(ConfigValidationStep):
                         m
                         for m in matches
                         if isinstance(
-                            candidate_conf := fv.get_declaring_config(m, result),
+                            candidate_conf := result.get_config_for_path(
+                                result.get_path_for_id(m)[:-1]
+                            ),
                             dict,
                         )
                         and all(

esphome/core/entity_helpers.py

@@ -646,12 +646,14 @@ def inherit_property_from(property_to_inherit, parent_id_property, transform=Non
 
             # Get config for the parent entity
             parent_id = _walk_config(config, parent_id_property)
-            parent_config = fv.get_declaring_config(parent_id, fconf)
+            parent_path = fconf.get_path_for_id(parent_id)[:-1]
+            parent_config = fconf.get_config_for_path(parent_path)
 
             # If parent sensor has the property set, inherit it
             if property in parent_config:
+                path = fconf.get_path_for_id(config[CONF_ID])[:-1]
                 this_config = _walk_config(
-                    fv.get_declaring_config(config[CONF_ID], fconf), property_path
+                    fconf.get_config_for_path(path), property_path
                 )
                 value = parent_config[property]
                 if transform:

esphome/pins.py

@@ -21,7 +21,6 @@ from esphome.const import (
 )
 from esphome.core import CORE, ID
 from esphome.cpp_generator import MockObj, MockObjClass
-import esphome.final_validate as fv
 from esphome.schema_extractors import schema_extractor
 from esphome.types import ConfigType, SafeExpType
 
@@ -134,7 +133,8 @@ class PinRegistry(dict):
                 with fconf.catch_error([cv.ROOT_CONFIG_PATH] + pin_path):
                     if final_val_fun is not None:
                         # Get the containing path of the config providing this pin.
-                        parent_config = fv.get_declaring_config(client_id, fconf)
+                        parent_path = fconf.get_path_for_id(client_id)[:-1]
+                        parent_config = fconf.get_config_for_path(parent_path)
                         final_val_fun(pin_config, parent_config)
                     allow_others = pin_config.get(CONF_ALLOW_OTHER_USES, False)
                     if count != 1 and not allow_others and not CORE.testing_mode:

tests/integration/state_utils.py

@@ -595,8 +595,8 @@ class SensorStateCollector:
 class SensorTracker:
     """Data-driven sensor state tracker with expected-value futures.
 
-    Tracks sensor and binary sensor state updates and resolves futures when
-    they report specific expected values. Eliminates per-sensor future
+    Tracks sensor, binary sensor and text sensor state updates and resolves
+    futures when they report specific expected values. Eliminates per-sensor future
     boilerplate.
 
     Usage::
@@ -608,12 +608,20 @@ class SensorTracker:
     """
 
     def __init__(self, sensor_names: list[str]) -> None:
-        self.sensor_states: dict[str, list[float]] = {name: [] for name in sensor_names}
+        self.sensor_states: dict[str, list[float | bool | str]] = {
+            name: [] for name in sensor_names
+        }
         self.key_to_sensor: dict[int, str] = {}
         self._expectations: dict[str, list[tuple[object, asyncio.Future]]] = {}
 
     _ANY = object()  # Sentinel: match any value
 
+    def map_entities(self, entities: list[EntityInfo]) -> None:
+        """Map the tracked names to the keys of ``entities``."""
+        self.key_to_sensor.update(
+            build_key_to_entity_mapping(entities, list(self.sensor_states))
+        )
+
     def expect(self, name: str, value: object) -> asyncio.Future:
         """Register an expected value for *name* and return a future for it."""
         future: asyncio.Future = asyncio.get_running_loop().create_future()
@@ -638,7 +646,7 @@ class SensorTracker:
                 connect-time states so they cannot satisfy a later phase.
         """
         if (
-            not isinstance(state, (SensorState, BinarySensorState))
+            not isinstance(state, (SensorState, BinarySensorState, TextSensorState))
             or state.missing_state
         ):
             return
@@ -708,9 +716,7 @@ class SensorTracker:
                 satisfy a later phase.
         """
         entities, _ = await client.list_entities_services()
-        self.key_to_sensor.update(
-            build_key_to_entity_mapping(entities, list(self.sensor_states.keys()))
-        )
+        self.map_entities(entities)
         initial_state_helper = InitialStateHelper(entities)
         client.subscribe_states(initial_state_helper.on_state_wrapper(self.on_state))
         await initial_state_helper.wait_for_initial_states()

tests/integration/test_api_get_time_response_timezone.py

@@ -6,7 +6,7 @@ from aioesphomeapi import connection as api_connection
 from aioesphomeapi.api_pb2 import GetTimeResponse
 import pytest
 
-from .state_utils import SensorTracker, build_key_to_entity_mapping
+from .state_utils import SensorTracker
 from .types import APIClientConnectedFactory, RunCompiledFunction
 
 # 2024-01-01 00:00:00 UTC
@@ -32,7 +32,7 @@ async def test_api_get_time_response_timezone(
     async with run_compiled(yaml_config), api_client_connected() as client:
         entities, _ = await client.list_entities_services()
         tracker = SensorTracker(["tz_offset"])
-        tracker.key_to_sensor = build_key_to_entity_mapping(entities, ["tz_offset"])
+        tracker.map_entities(entities)
         client.subscribe_states(tracker.on_state)
 
         await tracker.await_change(tracker.expect_any("tz_offset"), "tz_offset")

tests/integration/test_status_flags.py

@@ -19,11 +19,7 @@ import asyncio
 
 import pytest
 
-from .state_utils import (
-    SensorTracker,
-    build_key_to_entity_mapping,
-    subscribe_after_initial_states,
-)
+from .state_utils import SensorTracker, subscribe_after_initial_states
 from .types import APIClientConnectedFactory, RunCompiledFunction
 
 # Time to let the host-mode main loop run so status_led_light.loop() can
@@ -67,9 +63,7 @@ async def test_status_flags(
                 "status_led_last_state",
             ]
         )
-        tracker.key_to_sensor.update(
-            build_key_to_entity_mapping(entities, list(tracker.sensor_states.keys()))
-        )
+        tracker.map_entities(entities)
 
         # Swallow initial state broadcasts so the test only reacts to
         # state changes triggered by our service calls.

tests/integration/test_template_time.py

@@ -2,14 +2,10 @@
 
 from __future__ import annotations
 
-import asyncio
-
-from aioesphomeapi import EntityState, TextSensorState
 import pytest
 
-from .state_utils import build_key_to_entity_mapping, subscribe_after_initial_states
+from .state_utils import SensorTracker, subscribe_after_initial_states
 from .types import APIClientConnectedFactory, RunCompiledFunction
-from .wait_utils import wait_or_fail
 
 # Timestamps are exposed as text sensors (rather than 32-bit float sensor states,
 # which cannot represent a UNIX epoch exactly) so the exact value can be checked.
@@ -30,18 +26,10 @@ async def test_template_time(
     async with run_compiled(yaml_config), api_client_connected() as client:
         entities, _ = await client.list_entities_services()
 
-        key_to_name = build_key_to_entity_mapping(entities, list(EXPECTED_STATES))
-        events = {name: asyncio.Event() for name in EXPECTED_STATES}
-
-        def on_state(state: EntityState) -> None:
-            if isinstance(state, TextSensorState) and not state.missing_state:
-                name = key_to_name.get(state.key)
-                if name is not None and state.state == EXPECTED_STATES[name]:
-                    events[name].set()
-
-        await subscribe_after_initial_states(client, entities, on_state=on_state)
-
-        for name, event in events.items():
-            await wait_or_fail(
-                event, 3.0, f"Timeout waiting for {name} to report its expected value"
-            )
+        tracker = SensorTracker(list(EXPECTED_STATES))
+        tracker.map_entities(entities)
+        futures = tracker.expect_all(EXPECTED_STATES)
+        await subscribe_after_initial_states(
+            client, entities, on_state=tracker.on_state
+        )
+        await tracker.await_all(futures, timeout=3.0)