Merge remote-tracking branch 'origin/fix-water-heater-device-id' into integration

esphome/components/api/api_connection.cpp

@@ -1346,9 +1346,8 @@ uint16_t APIConnection::try_send_water_heater_state(EntityBase *entity, APIConne
   resp.target_temperature_low = wh->get_target_temperature_low();
   resp.target_temperature_high = wh->get_target_temperature_high();
   resp.state = wh->get_state();
-  resp.key = wh->get_object_id_hash();
 
-  return encode_message_to_buffer(resp, WaterHeaterStateResponse::MESSAGE_TYPE, conn, remaining_size);
+  return fill_and_encode_entity_state(wh, resp, WaterHeaterStateResponse::MESSAGE_TYPE, conn, remaining_size);
 }
 uint16_t APIConnection::try_send_water_heater_info(EntityBase *entity, APIConnection *conn, uint32_t remaining_size) {
   auto *wh = static_cast<water_heater::WaterHeater *>(entity);

tests/integration/fixtures/device_id_in_state.yaml

@@ -46,6 +46,7 @@ sensor:
 
 binary_sensor:
   - platform: template
+    id: motion_detected
     name: Motion Detected
     device_id: motion_sensor
     lambda: return true;
@@ -82,3 +83,117 @@ output:
     write_action:
       - lambda: |-
           ESP_LOGD("test", "Light output: %d", state);
+
+cover:
+  - platform: template
+    name: Garage Door
+    device_id: motion_sensor
+    optimistic: true
+
+fan:
+  - platform: template
+    name: Ceiling Fan
+    device_id: humidity_monitor
+    speed_count: 3
+    has_oscillating: false
+    has_direction: false
+
+lock:
+  - platform: template
+    name: Front Door Lock
+    device_id: motion_sensor
+    optimistic: true
+
+number:
+  - platform: template
+    name: Target Temperature
+    device_id: temperature_monitor
+    optimistic: true
+    min_value: 0
+    max_value: 100
+    step: 1
+
+select:
+  - platform: template
+    name: Mode Select
+    device_id: humidity_monitor
+    optimistic: true
+    options:
+      - "Auto"
+      - "Manual"
+
+text:
+  - platform: template
+    name: Device Label
+    device_id: temperature_monitor
+    optimistic: true
+    mode: text
+
+valve:
+  - platform: template
+    name: Water Valve
+    device_id: humidity_monitor
+    optimistic: true
+
+globals:
+  - id: global_away
+    type: bool
+    initial_value: "false"
+  - id: global_is_on
+    type: bool
+    initial_value: "true"
+
+water_heater:
+  - platform: template
+    name: Test Boiler
+    device_id: temperature_monitor
+    optimistic: true
+    current_temperature: !lambda "return 45.0f;"
+    target_temperature: !lambda "return 60.0f;"
+    away: !lambda "return id(global_away);"
+    is_on: !lambda "return id(global_is_on);"
+    supported_modes:
+      - "off"
+      - electric
+    visual:
+      min_temperature: 30.0
+      max_temperature: 85.0
+      target_temperature_step: 0.5
+    set_action:
+      - lambda: |-
+          ESP_LOGD("test", "Water heater set");
+
+alarm_control_panel:
+  - platform: template
+    name: House Alarm
+    device_id: motion_sensor
+    codes:
+      - "1234"
+    restore_mode: ALWAYS_DISARMED
+    binary_sensors:
+      - input: motion_detected
+
+datetime:
+  - platform: template
+    name: Schedule Date
+    device_id: temperature_monitor
+    type: date
+    optimistic: true
+  - platform: template
+    name: Schedule Time
+    device_id: humidity_monitor
+    type: time
+    optimistic: true
+  - platform: template
+    name: Schedule DateTime
+    device_id: motion_sensor
+    type: datetime
+    optimistic: true
+
+event:
+  - platform: template
+    name: Doorbell
+    device_id: motion_sensor
+    event_types:
+      - "press"
+      - "double_press"

tests/integration/test_device_id_in_state.py

@@ -4,11 +4,80 @@ from __future__ import annotations
 
 import asyncio
 
-from aioesphomeapi import BinarySensorState, EntityState, SensorState, TextSensorState
+from aioesphomeapi import (
+    AlarmControlPanelEntityState,
+    BinarySensorState,
+    CoverState,
+    DateState,
+    DateTimeState,
+    EntityState,
+    FanState,
+    LightState,
+    LockEntityState,
+    NumberState,
+    SelectState,
+    SensorState,
+    SwitchState,
+    TextSensorState,
+    TextState,
+    TimeState,
+    ValveState,
+    WaterHeaterState,
+)
 import pytest
 
 from .types import APIClientConnectedFactory, RunCompiledFunction
 
+# Mapping of entity name to device name for all entities with device_id
+ENTITY_TO_DEVICE = {
+    # Original entities
+    "Temperature": "Temperature Monitor",
+    "Humidity": "Humidity Monitor",
+    "Motion Detected": "Motion Sensor",
+    "Temperature Monitor Power": "Temperature Monitor",
+    "Temperature Status": "Temperature Monitor",
+    "Motion Light": "Motion Sensor",
+    # New entity types
+    "Garage Door": "Motion Sensor",
+    "Ceiling Fan": "Humidity Monitor",
+    "Front Door Lock": "Motion Sensor",
+    "Target Temperature": "Temperature Monitor",
+    "Mode Select": "Humidity Monitor",
+    "Device Label": "Temperature Monitor",
+    "Water Valve": "Humidity Monitor",
+    "Test Boiler": "Temperature Monitor",
+    "House Alarm": "Motion Sensor",
+    "Schedule Date": "Temperature Monitor",
+    "Schedule Time": "Humidity Monitor",
+    "Schedule DateTime": "Motion Sensor",
+    "Doorbell": "Motion Sensor",
+}
+
+# Entities without device_id (should have device_id 0)
+NO_DEVICE_ENTITIES = {"No Device Sensor"}
+
+# State types that should have non-zero device_id, mapped by their aioesphomeapi class
+EXPECTED_STATE_TYPES = [
+    (SensorState, "sensor"),
+    (BinarySensorState, "binary_sensor"),
+    (SwitchState, "switch"),
+    (TextSensorState, "text_sensor"),
+    (LightState, "light"),
+    (CoverState, "cover"),
+    (FanState, "fan"),
+    (LockEntityState, "lock"),
+    (NumberState, "number"),
+    (SelectState, "select"),
+    (TextState, "text"),
+    (ValveState, "valve"),
+    (WaterHeaterState, "water_heater"),
+    (AlarmControlPanelEntityState, "alarm_control_panel"),
+    (DateState, "date"),
+    (TimeState, "time"),
+    (DateTimeState, "datetime"),
+    # Event is stateless (no initial state sent on subscribe)
+]
+
 
 @pytest.mark.asyncio
 async def test_device_id_in_state(
@@ -40,34 +109,35 @@ async def test_device_id_in_state(
         entity_device_mapping: dict[int, int] = {}
 
         for entity in all_entities:
-            # All entities have name and key attributes
-            if entity.name == "Temperature":
-                entity_device_mapping[entity.key] = device_ids["Temperature Monitor"]
-            elif entity.name == "Humidity":
-                entity_device_mapping[entity.key] = device_ids["Humidity Monitor"]
-            elif entity.name == "Motion Detected":
-                entity_device_mapping[entity.key] = device_ids["Motion Sensor"]
-            elif entity.name in {"Temperature Monitor Power", "Temperature Status"}:
-                entity_device_mapping[entity.key] = device_ids["Temperature Monitor"]
-            elif entity.name == "Motion Light":
-                entity_device_mapping[entity.key] = device_ids["Motion Sensor"]
-            elif entity.name == "No Device Sensor":
-                # Entity without device_id should have device_id 0
+            if entity.name in ENTITY_TO_DEVICE:
+                expected_device = ENTITY_TO_DEVICE[entity.name]
+                entity_device_mapping[entity.key] = device_ids[expected_device]
+            elif entity.name in NO_DEVICE_ENTITIES:
                 entity_device_mapping[entity.key] = 0
 
-        assert len(entity_device_mapping) >= 6, (
-            f"Expected at least 6 mapped entities, got {len(entity_device_mapping)}"
+        expected_count = len(ENTITY_TO_DEVICE) + len(NO_DEVICE_ENTITIES)
+        assert len(entity_device_mapping) >= expected_count, (
+            f"Expected at least {expected_count} mapped entities, "
+            f"got {len(entity_device_mapping)}. "
+            f"Missing: {set(ENTITY_TO_DEVICE) | NO_DEVICE_ENTITIES - {e.name for e in all_entities}}"
+        )
+
+        # Subscribe to states and wait for all mapped entities
+        # Event entities are stateless (no initial state on subscribe),
+        # so exclude them from the expected count
+        stateless_keys = {e.key for e in all_entities if e.name == "Doorbell"}
+        stateful_count = len(entity_device_mapping) - len(
+            stateless_keys & entity_device_mapping.keys()
         )
 
-        # Subscribe to states
         loop = asyncio.get_running_loop()
         states: dict[int, EntityState] = {}
         states_future: asyncio.Future[bool] = loop.create_future()
 
         def on_state(state: EntityState) -> None:
-            states[state.key] = state
-            # Check if we have states for all mapped entities
-            if len(states) >= len(entity_device_mapping) and not states_future.done():
+            if state.key in entity_device_mapping:
+                states[state.key] = state
+            if len(states) >= stateful_count and not states_future.done():
                 states_future.set_result(True)
 
         client.subscribe_states(on_state)
@@ -76,9 +146,16 @@ async def test_device_id_in_state(
         try:
             await asyncio.wait_for(states_future, timeout=10.0)
         except TimeoutError:
+            received_names = {e.name for e in all_entities if e.key in states}
+            missing_names = (
+                (set(ENTITY_TO_DEVICE) | NO_DEVICE_ENTITIES)
+                - received_names
+                - {"Doorbell"}
+            )
             pytest.fail(
                 f"Did not receive all entity states within 10 seconds. "
-                f"Received {len(states)} states, expected {len(entity_device_mapping)}"
+                f"Received {len(states)} states. "
+                f"Missing: {missing_names}"
             )
 
         # Verify each state has the correct device_id
@@ -86,51 +163,33 @@ async def test_device_id_in_state(
         for key, expected_device_id in entity_device_mapping.items():
             if key in states:
                 state = states[key]
+                entity_name = next(
+                    (e.name for e in all_entities if e.key == key), f"key={key}"
+                )
 
                 assert state.device_id == expected_device_id, (
-                    f"State for key {key} has device_id {state.device_id}, "
-                    f"expected {expected_device_id}"
+                    f"State for '{entity_name}' (type={type(state).__name__}) "
+                    f"has device_id {state.device_id}, expected {expected_device_id}"
                 )
                 verified_count += 1
 
-        assert verified_count >= 6, (
-            f"Only verified {verified_count} states, expected at least 6"
+        # All stateful entities should be verified (everything except Doorbell event)
+        expected_verified = expected_count - 1  # exclude Doorbell
+        assert verified_count >= expected_verified, (
+            f"Only verified {verified_count} states, expected at least {expected_verified}"
         )
 
-        # Test specific state types to ensure device_id is present
-        # Find a sensor state with device_id
-        sensor_state = next(
-            (
+        # Verify each expected state type has at least one instance with non-zero device_id
+        for state_type, type_name in EXPECTED_STATE_TYPES:
+            matching = [
                 s
                 for s in states.values()
-                if isinstance(s, SensorState)
-                and isinstance(s.state, float)
-                and s.device_id != 0
-            ),
-            None,
-        )
-        assert sensor_state is not None, "No sensor state with device_id found"
-        assert sensor_state.device_id > 0, "Sensor state should have non-zero device_id"
-
-        # Find a binary sensor state
-        binary_sensor_state = next(
-            (s for s in states.values() if isinstance(s, BinarySensorState)),
-            None,
-        )
-        assert binary_sensor_state is not None, "No binary sensor state found"
-        assert binary_sensor_state.device_id > 0, (
-            "Binary sensor state should have non-zero device_id"
-        )
-
-        # Find a text sensor state
-        text_sensor_state = next(
-            (s for s in states.values() if isinstance(s, TextSensorState)),
-            None,
-        )
-        assert text_sensor_state is not None, "No text sensor state found"
-        assert text_sensor_state.device_id > 0, (
-            "Text sensor state should have non-zero device_id"
-        )
+                if isinstance(s, state_type) and s.device_id != 0
+            ]
+            assert matching, (
+                f"No {type_name} state (type={state_type.__name__}) "
+                f"with non-zero device_id found"
+            )
 
         # Verify the "No Device Sensor" has device_id = 0
         no_device_key = next(