[tests] Fix flaky modbus server/controller integration tests (#18258)

tests/integration/state_utils.py

@@ -420,8 +420,15 @@ class SensorTracker:
         """Call ``expect`` for every entry and return a dict of futures."""
         return {name: self.expect(name, value) for name, value in expected.items()}
 
-    def on_state(self, state: EntityState) -> None:
-        """State callback suitable for ``subscribe_states``."""
+    def on_state(self, state: EntityState, first_pending_only: bool = False) -> None:
+        """State callback suitable for ``subscribe_states``.
+
+        Args:
+            state: The state update to record
+            first_pending_only: Only allow the first pending expectation for this
+                sensor to match, instead of the first matching one. Used for
+                connect-time states so they cannot satisfy a later phase.
+        """
         if (
             not isinstance(state, (SensorState, BinarySensorState))
             or state.missing_state
@@ -432,11 +439,13 @@ class SensorTracker:
             return
         self.sensor_states[sensor_name].append(state.state)
         for expected_value, future in self._expectations.get(sensor_name, []):
-            if not future.done() and (
-                expected_value is self._ANY or state.state == expected_value
-            ):
+            if future.done():
+                continue
+            if expected_value is self._ANY or state.state == expected_value:
                 future.set_result(True)
                 break
+            if first_pending_only:
+                break
 
     async def await_change(
         self, future: asyncio.Future, name: str, timeout: float = 2.0
@@ -474,8 +483,22 @@ class SensorTracker:
         for name, future in futures.items():
             await self.await_change(future, name, timeout=timeout)
 
-    async def setup_and_start_scenario(self, client) -> list:
-        """Wire up subscriptions, wait for initial states, press Start Scenario."""
+    async def setup_and_start_scenario(
+        self, client: APIClient, match_initial_states: bool = False
+    ) -> list[EntityInfo]:
+        """Wire up subscriptions, wait for initial states, press Start Scenario.
+
+        Args:
+            client: The connected API client
+            match_initial_states: Also match expectations against the states the
+                device sends when the client connects, so a value published before
+                the client subscribed still counts. Binary sensors need this: they
+                drop repeats, so a value that lands in the connect-time dump is
+                never sent again. Plain sensors publish on every poll, so there it
+                only saves waiting for the next one. Only the first pending
+                expectation per sensor can match, so a connect-time value cannot
+                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()))
@@ -488,6 +511,9 @@ class SensorTracker:
             import pytest
 
             pytest.fail("Timeout waiting for initial states")
+        if match_initial_states:
+            for state in initial_state_helper.initial_states.values():
+                self.on_state(state, first_pending_only=True)
         start_btn = find_entity(entities, "start_scenario", ButtonInfo)
         assert start_btn is not None, "Start Scenario button not found"
         client.button_command(start_btn.key)

tests/integration/test_uart_mock_modbus.py

@@ -330,7 +330,10 @@ async def test_uart_mock_modbus_server_controller(
         run_compiled(yaml_config, line_callback=line_callback),
         api_client_connected() as client,
     ):
-        await tracker.setup_and_start_scenario(client)
+        # The controller polls from boot, so the first values can already be in
+        # the states the device sends on connect; matching them there saves
+        # waiting for the next poll
+        await tracker.setup_and_start_scenario(client, match_initial_states=True)
         await tracker.await_all(futures)
         _assert_no_modbus_errors(error_log_lines, warning_log_lines)
 
@@ -392,7 +395,12 @@ async def test_uart_mock_modbus_server_controller_write(
         run_compiled(yaml_config, line_callback=line_callback),
         api_client_connected() as client,
     ):
-        entities = await tracker.setup_and_start_scenario(client)
+        # The controller polls from boot, so the baseline can already be in the
+        # states the device sends on connect; matching it there saves waiting for
+        # the next poll
+        entities = await tracker.setup_and_start_scenario(
+            client, match_initial_states=True
+        )
 
         # Wait for initial baseline values to confirm the controller <-> server
         # connection is working before issuing writes
@@ -456,7 +464,11 @@ async def test_uart_mock_modbus_server_controller_bits(
         run_compiled(yaml_config, line_callback=line_callback),
         api_client_connected() as client,
     ):
-        entities = await tracker.setup_and_start_scenario(client)
+        # The controller polls from boot and binary sensors drop repeats, so the
+        # baseline can arrive only in the states the device sends on connect
+        entities = await tracker.setup_and_start_scenario(
+            client, match_initial_states=True
+        )
 
         # Wait for initial baseline values to confirm the controller <-> server
         # connection is working before issuing writes
@@ -491,7 +503,10 @@ async def test_uart_mock_modbus_server_controller_multiple(
         run_compiled(yaml_config, line_callback=line_callback),
         api_client_connected() as client,
     ):
-        await tracker.setup_and_start_scenario(client)
+        # The controller polls from boot, so the first values can already be in
+        # the states the device sends on connect; matching them there saves
+        # waiting for the next poll
+        await tracker.setup_and_start_scenario(client, match_initial_states=True)
         await tracker.await_all(futures)
         _assert_no_modbus_errors(error_log_lines, warning_log_lines)