[Fix] Websocket and Log Viewer

tests/test_run/test_websocket_socket.py

@@ -28,7 +28,14 @@
     TestSuiteExecution,
     TestSuiteMetadata,
 )
-from th_cli.test_run.socket_schemas import TestCaseUpdate, TestRunUpdate, TestStepUpdate, TestSuiteUpdate, TestUpdate
+from th_cli.test_run.socket_schemas import (
+    TestCaseUpdate,
+    TestLogRecord,
+    TestRunUpdate,
+    TestStepUpdate,
+    TestSuiteUpdate,
+    TestUpdate,
+)
 from th_cli.test_run.websocket import TestRunSocket
 
 # ---------------------------------------------------------------------------
@@ -304,7 +311,7 @@ async def test_step_update_routed_correctly(self):
             ),
         )
         with patch.object(s, "_TestRunSocket__log_test_step_update") as mock_fn:
-            await s._TestRunSocket__handle_test_update(socket=AsyncMock(), update=update)
+            await s._TestRunSocket__handle_test_update(update=update)
 
         mock_fn.assert_called_once()
 
@@ -319,7 +326,7 @@ async def test_case_update_routed_correctly(self):
             body=TestCaseUpdate(state="passed", test_case_execution_index=0, test_suite_execution_index=0),
         )
         with patch.object(s, "_TestRunSocket__log_test_case_update") as mock_fn:
-            await s._TestRunSocket__handle_test_update(socket=AsyncMock(), update=update)
+            await s._TestRunSocket__handle_test_update(update=update)
 
         mock_fn.assert_called_once()
 
@@ -333,28 +340,55 @@ async def test_suite_update_routed_correctly(self):
             body=TestSuiteUpdate(state="passed", test_suite_execution_index=0),
         )
         with patch.object(s, "_TestRunSocket__log_test_suite_update") as mock_fn:
-            await s._TestRunSocket__handle_test_update(socket=AsyncMock(), update=update)
+            await s._TestRunSocket__handle_test_update(update=update)
 
         mock_fn.assert_called_once()
 
     @pytest.mark.asyncio
-    async def test_run_update_executing_does_not_close_socket(self):
+    async def test_run_update_executing_leaves_run_not_finished(self):
         s = _make_socket()
-        mock_socket = AsyncMock()
 
         update = TestUpdate(test_type="test_run", body=TestRunUpdate(state="executing", test_run_execution_id=1))
         with patch.object(s, "_TestRunSocket__log_test_run_update", new_callable=AsyncMock):
-            await s._TestRunSocket__handle_test_update(socket=mock_socket, update=update)
+            await s._TestRunSocket__handle_test_update(update=update)
 
-        mock_socket.close.assert_not_called()
+        assert s._run_finished is False
 
     @pytest.mark.asyncio
-    async def test_run_update_non_executing_closes_socket(self):
+    async def test_run_update_non_executing_marks_run_finished(self):
         s = _make_socket()
-        mock_socket = AsyncMock()
 
         update = TestUpdate(test_type="test_run", body=TestRunUpdate(state="passed", test_run_execution_id=1))
         with patch.object(s, "_TestRunSocket__log_test_run_update", new_callable=AsyncMock):
-            await s._TestRunSocket__handle_test_update(socket=mock_socket, update=update)
+            await s._TestRunSocket__handle_test_update(update=update)
+
+        assert s._run_finished is True
+
+    @pytest.mark.asyncio
+    async def test_run_update_pending_leaves_run_not_finished(self):
+        # Regression test: "pending" is non-terminal (backend's TestRun.completed()
+        # excludes both PENDING and EXECUTING), so it must not close the socket.
+        # A prior implementation used a negation check (`state != "executing"`)
+        # that misclassified any non-"executing" state, including "pending", as
+        # terminal.
+        s = _make_socket()
+
+        update = TestUpdate(test_type="test_run", body=TestRunUpdate(state="pending", test_run_execution_id=1))
+        with patch.object(s, "_TestRunSocket__log_test_run_update", new_callable=AsyncMock):
+            await s._TestRunSocket__handle_test_update(update=update)
+
+        assert s._run_finished is False
+
+    @pytest.mark.asyncio
+    async def test_handle_log_record_logs_every_record(self):
+        s = _make_socket()
+        records = [
+            TestLogRecord(level="INFO", timestamp=0.0, message=f"msg{i}") for i in range(3)
+        ]
+
+        with patch("th_cli.test_run.websocket.logger") as mock_logger:
+            await s._TestRunSocket__handle_log_record(records)
 
-        mock_socket.close.assert_called_once()
+        assert mock_logger.log.call_count == 3
+        for record in records:
+            mock_logger.log.assert_any_call(record.level, record.message)

th_cli/commands/run_tests.py

@@ -224,6 +224,7 @@ async def run_tests(
             execution_pics=execution_pics,
             project_id=project_id,
         )
+        test_logging.set_download_run_id(new_test_run.id)
         if _contains_webrtc_two_way_talk(selected_tests_dict):
             _webrtc_handler = TwoWayTalkHandler(port=8999)
             _webrtc_handler.start_waiting()

th_cli/test_run/log_stream_handler.py

@@ -28,56 +28,72 @@ class LogStreamHandler:
 
     def __init__(self, port: int = 8998):
         """Initialize the log stream handler.
-        
+
         Args:
             port: Port number for the HTTP server (default: 8998)
         """
         self.port = port
         self.http_server = LogsHTTPServer(port=port)
         self.log_queue: queue.Queue = queue.Queue(maxsize=1000)
         self.is_running = False
-        self.log_file_path: Optional[str] = None
-
-    def start(self, test_run_title: str = "Test Execution", log_file_path: Optional[str] = None) -> str:
+
+    def start(self, test_run_title: str = "Test Execution") -> str:
         """Start the log streaming HTTP server.
-        
+
         Args:
             test_run_title: Title of the test run for display
-            log_file_path: Path to the log file for download functionality
-
+
         Returns:
             URL where logs can be viewed
         """
         if self.is_running:
             logger.warning("Log stream handler already running")
             return self._get_log_viewer_url()
-        
+
         try:
-            # Store log file path for download functionality
-            self.log_file_path = log_file_path
-
             # Get local IP address
             local_ip = self._get_local_ip()
-            
+
             # Start HTTP server
             self.http_server.start(
                 log_queue=self.log_queue,
                 test_run_title=test_run_title,
                 local_ip=local_ip,
-                log_file_path=log_file_path,
             )
-            
+
             self.is_running = True
-            
+
             viewer_url = f"http://{local_ip}:{self.port}"
             logger.info(f"Log stream viewer started: {viewer_url}")
-            
+
             return viewer_url
-            
+
         except Exception as e:
             logger.error(f"Failed to start log stream handler: {e}")
             raise
-
+
+    def set_run_id(self, run_id: int) -> None:
+        """Tell the HTTP server which run's log to link "Download Logs" to,
+        once the run has been created (its id isn't known when the server
+        starts).
+        """
+        if not self.is_running:
+            return
+
+        self.http_server.set_run_id(run_id)
+
+        # A viewer may already be connected (the run_id is typically set
+        # only *after* the viewer URL was printed and likely opened), so
+        # also push it through the existing SSE stream as a control message
+        # - a future/refreshed page load will pick it up from the HTTP
+        # server attribute above, but an already-open one only sees this.
+        try:
+            self.log_queue.put_nowait({"__event__": "run_id", "run_id": run_id})
+        except queue.Full:
+            # Best-effort: a future page load/refresh will still pick up
+            # the run id via the HTTP server attribute set above.
+            pass
+
     def stop(self):
         """Stop the log streaming HTTP server."""
         if not self.is_running: