From 72bfcc1c3024b7c731b83a0460b4e393cc057574 Mon Sep 17 00:00:00 2001 From: Steven Green Date: Fri, 2 Oct 2026 12:10:47 -0700 Subject: [PATCH] [Feature] End test runs faster Currently, the CLI waits 5s after a test run reports it is complete, so that the log records the backend flushes after its terminal state update still reach the local log file. This is usually extremely generous. Backends that send the terminal test run update after the run's last log records now mark it with `logs_complete`. When that field is present, the CLI closes the socket as soon as the update arrives instead of draining, saving about 4s/run. If it is false, the run ended before its logs could be flushed, and the CLI warns that the log file may be incomplete. Older backends don't send the field, so the CLI keeps the 5s drain for them. --- tests/test_run/test_websocket_connect.py | 69 ++++++++++++++++++++++++ th_cli/test_run/socket_schemas.py | 5 ++ th_cli/test_run/websocket.py | 35 ++++++++---- 3 files changed, 98 insertions(+), 11 deletions(-) diff --git a/tests/test_run/test_websocket_connect.py b/tests/test_run/test_websocket_connect.py index bea5e59..4a1a8b3 100644 --- a/tests/test_run/test_websocket_connect.py +++ b/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() diff --git a/th_cli/test_run/socket_schemas.py b/th_cli/test_run/socket_schemas.py index 828fcfc..d8ad65c 100644 --- a/th_cli/test_run/socket_schemas.py +++ b/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): diff --git a/th_cli/test_run/websocket.py b/th_cli/test_run/websocket.py index d6e1b98..9ab33c7 100644 --- a/th_cli/test_run/websocket.py +++ b/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)