[Feature] End test runs faster

tests/test_run/test_websocket_connect.py

@@ -20,6 +20,8 @@
 silently treated as a successfully completed run.
 """
 
+import asyncio
+import json
 from unittest.mock import AsyncMock, patch
 
 import pytest
@@ -166,3 +168,70 @@ async def test_abrupt_close_after_run_finished_does_not_raise(self):
             await s.connect_websocket()  # must not raise
 
         fake_socket.close.assert_called_once()
+
+
+# ---------------------------------------------------------------------------
+# connect_websocket — closing after the terminal TestRunUpdate
+# ---------------------------------------------------------------------------
+
+
+def _terminal_run_update(**body) -> str:
+    return json.dumps(
+        {
+            "type": "test_update",
+            "payload": {
+                "test_type": "test_run",
+                "body": {"state": "passed", "test_run_execution_id": 1, **body},
+            },
+        }
+    )
+
+
+_TRAILING_LOG_RECORDS = json.dumps(
+    {"type": "test_log_records", "payload": [{"level": "INFO", "timestamp": 1.0, "message": "trailing"}]}
+)
+
+
+@pytest.mark.unit
+class TestConnectWebsocketAfterTerminalUpdate:
+    @pytest.mark.asyncio
+    async def test_logs_complete_closes_without_draining(self):
+        """The backend sent the terminal update after the run's last log records."""
+        s = _make_socket()
+        fake_socket = _FakeWSSocket(
+            recv_side_effect=[_terminal_run_update(logs_complete=True), AssertionError("must not drain")]
+        )
+
+        with _patch_connect(fake_socket), patch("th_cli.test_run.websocket.DRAIN_TIMEOUT_S", 3600):
+            await asyncio.wait_for(s.connect_websocket(), timeout=5)
+
+        assert fake_socket.recv.await_count == 1
+        fake_socket.close.assert_called_once()
+
+    @pytest.mark.asyncio
+    async def test_logs_incomplete_warns_and_closes_without_draining(self, capsys):
+        s = _make_socket()
+        fake_socket = _FakeWSSocket(
+            recv_side_effect=[_terminal_run_update(logs_complete=False), AssertionError("must not drain")]
+        )
+
+        with _patch_connect(fake_socket), patch("th_cli.test_run.websocket.DRAIN_TIMEOUT_S", 3600):
+            await asyncio.wait_for(s.connect_websocket(), timeout=5)
+
+        assert fake_socket.recv.await_count == 1
+        assert "couldn't send all of the run's log records" in capsys.readouterr().err
+
+    @pytest.mark.asyncio
+    async def test_older_backend_drains_trailing_log_records(self):
+        """No logs_complete field: the backend's last log records may still follow."""
+        s = _make_socket()
+        fake_socket = _FakeWSSocket(
+            recv_side_effect=[_terminal_run_update(), _TRAILING_LOG_RECORDS, asyncio.TimeoutError()]
+        )
+
+        with _patch_connect(fake_socket), patch("th_cli.test_run.websocket.logger") as mock_logger:
+            await s.connect_websocket()
+
+        mock_logger.log.assert_any_call("INFO", "trailing")
+        assert fake_socket.recv.await_count == 3
+        fake_socket.close.assert_called_once()

th_cli/test_run/socket_schemas.py

@@ -40,6 +40,11 @@ class TestUpdateBase(BaseModel):
 
 class TestRunUpdate(TestUpdateBase):
     test_run_execution_id: int
+    # Set on the terminal update by backends that send it after the run's last
+    # log records: True if they were all sent, False if the run ended before
+    # its logs could be flushed. None from older backends, which send the
+    # terminal update before their last log records.
+    logs_complete: bool | None = None
 
 
 class TestSuiteUpdate(TestUpdateBase):

th_cli/test_run/websocket.py

@@ -38,6 +38,7 @@
     colorize_hierarchy_prefix,
     colorize_key_value,
     colorize_state,
+    colorize_warning,
 )
 from th_cli.config import config
 from th_cli.shared_constants import MessageTypeEnum, TestStateEnum
@@ -84,12 +85,12 @@
     (TestStateEnum.CANCELLED.value, "cancelled"),
 ]
 
-# After the test run reaches a terminal state, the backend may still have a
-# trailing batch of log records queued/in-flight (it flushes and broadcasts
-# any pending log entries *after* sending the terminal state update - see
-# TestLogHandler.finish()/TestUIObserver.complete_tasks() on the backend).
-# Keep draining for a short grace period instead of closing immediately, so
-# that trailing batch isn't dropped by a socket we already hung up on.
+# Older backends send the terminal test run update before flushing and
+# broadcasting the run's last log records. Their terminal update has no
+# `logs_complete` field; keep draining for a short grace period after it
+# instead of closing immediately, so that trailing batch isn't dropped by a
+# socket we already hung up on. Newer backends send the terminal update last,
+# so the socket closes as soon as it arrives.
 DRAIN_TIMEOUT_S = 5.0
 
 # Yield to the event loop every N log records while processing one batch, so
@@ -128,6 +129,8 @@ def __init__(
         self.two_way_talk_handler = two_way_talk_handler
         self._chip_server_info_displayed = False
         self._run_finished = False
+        # From the terminal TestRunUpdate (see TestRunUpdate.logs_complete).
+        self._logs_complete: bool | None = None
         # Track test step errors for logging
         # Key: (suite_index, case_index), Value: list of error strings from all steps
         self.test_case_step_errors: dict[tuple[int, int], list[str]] = {}
@@ -178,6 +181,10 @@ async def connect_websocket(self) -> None:
             ) as socket:
                 try:
                     while True:
+                        if self._run_finished and self._logs_complete is not None:
+                            # The backend sent the terminal update after the
+                            # run's last log records: nothing left to drain.
+                            break
                         try:
                             if self._run_finished:
                                 # Drain any trailing messages for a short grace
@@ -274,12 +281,18 @@ async def __handle_test_update(self, update: TestUpdate) -> None:
         elif isinstance(update.body, TestRunUpdate):
             await self.__log_test_run_update(update.body)
             if update.body.state not in NON_TERMINAL_RUN_STATES:
-                # Test run ended. Don't close immediately - the backend may
-                # still be flushing/broadcasting a trailing batch of log
-                # entries after this message; let the read loop keep
-                # draining for a short grace period (see DRAIN_TIMEOUT_S)
-                # before actually closing.
+                # Test run ended. The read loop closes the socket, after a
+                # drain period for older backends (see DRAIN_TIMEOUT_S).
                 self._run_finished = True
+                self._logs_complete = update.body.logs_complete
+                if self._logs_complete is False:
+                    click.echo(
+                        colorize_warning(
+                            "The backend couldn't send all of the run's log records; "
+                            "the log file may be missing its final lines."
+                        ),
+                        err=True,
+                    )
 
     async def __log_test_run_update(self, update: TestRunUpdate) -> None:
         # Display CHIP server info when test run starts executing (SDK container already running)