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)