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

Moved to its own pull request.

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)