| 1 | # Copyright 2026 Google LLC |
| 2 | # |
| 3 | # Licensed under the Apache License, Version 2.0 (the "License"); |
| 4 | # you may not use this file except in compliance with the License. |
| 5 | # You may obtain a copy of the License at |
| 6 | # |
| 7 | # http://www.apache.org/licenses/LICENSE-2.0 |
| 8 | # |
| 9 | # Unless required by applicable law or agreed to in writing, software |
| 10 | # distributed under the License is distributed on an "AS IS" BASIS, |
| 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 12 | # See the License for the specific language governing permissions and |
| 13 | # limitations under the License. |
| 14 | |
| 15 | import json |
| 16 | import os |
| 17 | import sys |
| 18 | import termios |
| 19 | from unittest.mock import MagicMock, patch |
| 20 | |
| 21 | from colab_cli.console import connect_console, on_message, on_open |
| 22 | from colab_cli.state import SessionState |
| 23 | import pytest |
| 24 | |
| 25 | |
| 26 | @pytest.fixture |
| 27 | def mock_session(): |
| 28 | return SessionState( |
| 29 | name="test-session", |
| 30 | token="test-token", |
| 31 | url="https://8080-m-s-kkb-usc1f1.us-central1-1.colab.dev", |
| 32 | endpoint="some-endpoint", |
| 33 | ) |
| 34 | |
| 35 | |
| 36 | @patch("colab_cli.console.websocket.WebSocketApp") |
| 37 | @patch("colab_cli.console.tty.setraw") |
| 38 | @patch("colab_cli.console.termios.tcgetattr") |
| 39 | @patch("colab_cli.console.termios.tcsetattr") |
| 40 | @patch("colab_cli.console.os.get_terminal_size") |
| 41 | @patch("colab_cli.console.sys.stdin.fileno") |
| 42 | @patch("colab_cli.console.sys.stdin.isatty") |
| 43 | def test_console_initialization( |
| 44 | mock_isatty, |
| 45 | mock_fileno, |
| 46 | mock_get_term_size, |
| 47 | mock_tcsetattr, |
| 48 | mock_tcgetattr, |
| 49 | mock_setraw, |
| 50 | mock_ws_app, |
| 51 | mock_session, |
| 52 | ): |
| 53 | # Setup mocks |
| 54 | mock_isatty.return_value = True |
| 55 | mock_fileno.return_value = 0 |
| 56 | mock_get_term_size.return_value = os.terminal_size((80, 24)) |
| 57 | mock_tcgetattr.return_value = ["fake_attrs"] |
| 58 | mock_ws_instance = MagicMock() |
| 59 | mock_ws_app.return_value = mock_ws_instance |
| 60 | |
| 61 | # We don't want run_forever to actually block or start threads in the test |
| 62 | mock_ws_instance.run_forever.return_value = None |
| 63 | |
| 64 | with patch("colab_cli.console.threading.Thread"): |
| 65 | connect_console(mock_session) |
| 66 | |
| 67 | # 1. Verify URL transformation |
| 68 | expected_url = "wss://8080-m-s-kkb-usc1f1.us-central1-1.colab.dev/colab/tty?colab-runtime-proxy-token=test-token" |
| 69 | mock_ws_app.assert_called_once() |
| 70 | assert mock_ws_app.call_args[1]["url"] == expected_url |
| 71 | |
| 72 | # 2. Verify raw mode setup and teardown |
| 73 | mock_tcgetattr.assert_called_once_with(sys.stdin.fileno()) |
| 74 | mock_setraw.assert_called_once_with(sys.stdin.fileno(), termios.TCSANOW) |
| 75 | |
| 76 | # Teardown should happen in a finally block |
| 77 | mock_tcsetattr.assert_called_once_with( |
| 78 | sys.stdin.fileno(), termios.TCSANOW, ["fake_attrs"] |
| 79 | ) |
| 80 | |
| 81 | |
| 82 | @patch("colab_cli.console.websocket.WebSocketApp") |
| 83 | @patch("colab_cli.console.tty.setraw") |
| 84 | @patch("colab_cli.console.termios.tcgetattr") |
| 85 | @patch("colab_cli.console.termios.tcsetattr") |
| 86 | @patch("colab_cli.console.sys.stdin.isatty") |
| 87 | def test_console_piped_input( |
| 88 | mock_isatty, |
| 89 | mock_tcsetattr, |
| 90 | mock_tcgetattr, |
| 91 | mock_setraw, |
| 92 | mock_ws_app, |
| 93 | mock_session, |
| 94 | ): |
| 95 | mock_isatty.return_value = False |
| 96 | mock_ws_instance = MagicMock() |
| 97 | mock_ws_app.return_value = mock_ws_instance |
| 98 | mock_ws_instance.run_forever.return_value = None |
| 99 | |
| 100 | with patch("colab_cli.console.threading.Thread"): |
| 101 | connect_console(mock_session) |
| 102 | |
| 103 | # In a piped environment, we should not attempt to use termios or tty |
| 104 | mock_tcgetattr.assert_not_called() |
| 105 | mock_setraw.assert_not_called() |
| 106 | mock_tcsetattr.assert_not_called() |
| 107 | |
| 108 | |
| 109 | @patch("colab_cli.console.os.get_terminal_size") |
| 110 | def test_on_open_sends_terminal_size(mock_get_term_size): |
| 111 | mock_ws = MagicMock() |
| 112 | mock_get_term_size.return_value = os.terminal_size((100, 40)) |
| 113 | |
| 114 | on_open(mock_ws) |
| 115 | |
| 116 | # Verify that the initial terminal size is sent |
| 117 | mock_ws.send.assert_called_once() |
| 118 | payload = json.loads(mock_ws.send.call_args[0][0]) |
| 119 | assert payload == {"cols": 100, "rows": 40} |
| 120 | |
| 121 | |
| 122 | @patch("colab_cli.console.sys.stdout.buffer.write") |
| 123 | @patch("colab_cli.console.sys.stdout.buffer.flush") |
| 124 | def test_on_message_writes_to_stdout(mock_flush, mock_write): |
| 125 | mock_ws = MagicMock() |
| 126 | test_data = "Hello \x1b[34mWorld\x1b[0m" |
| 127 | message_json = json.dumps({"data": test_data}) |
| 128 | |
| 129 | on_message(mock_ws, message_json) |
| 130 | |
| 131 | # Verify that the data is written exactly as received |
| 132 | mock_write.assert_called_once_with(test_data.encode("utf-8")) |
| 133 | mock_flush.assert_called_once() |
| 134 | |
| 135 | |
| 136 | @patch("colab_cli.console.os.get_terminal_size") |
| 137 | @patch("colab_cli.console.sys.stdin.isatty") |
| 138 | @patch("colab_cli.console.sys.stdin") |
| 139 | def test_read_stdin_eof_piped_sends_exit_and_closes_ws( |
| 140 | mock_stdin, mock_isatty, mock_get_term_size |
| 141 | ): |
| 142 | """When stdin is piped and reaches EOF, the read thread should send 'exit\\n' |
| 143 | to the remote shell and then close the websocket from the client side. |
| 144 | |
| 145 | The remote shell at /colab/tty is wrapped in tmux which swallows the bare |
| 146 | \\x04 (Ctrl-D) we used to send, so EOF used to leave the websocket open |
| 147 | indefinitely. Sending 'exit\\n' + ws.close() guarantees clean termination. |
| 148 | """ |
| 149 | import colab_cli.console as console_mod |
| 150 | |
| 151 | mock_isatty.return_value = False |
| 152 | # Simulate piped stdin: returns one line then EOF |
| 153 | mock_stdin.read.side_effect = ["e", "c", "h", "o", " ", "h", "i", "\n", ""] |
| 154 | mock_get_term_size.return_value = os.terminal_size((80, 24)) |
| 155 | |
| 156 | mock_ws = MagicMock() |
| 157 | |
| 158 | # on_open spawns the read thread; we want it to run synchronously here |
| 159 | # so we patch threading.Thread to call target immediately and join(). |
| 160 | real_thread = [] |
| 161 | |
| 162 | class SyncThread: |
| 163 | def __init__(self, target, daemon=None): |
| 164 | self.target = target |
| 165 | real_thread.append(self) |
| 166 | |
| 167 | def start(self): |
| 168 | self.target() |
| 169 | |
| 170 | console_mod._is_running = True |
| 171 | with patch("colab_cli.console.threading.Thread", SyncThread): |
| 172 | # Use a tiny grace period for the test |
| 173 | with patch("colab_cli.console.PIPED_EOF_GRACE_SECONDS", 0.01): |
| 174 | on_open(mock_ws) |
| 175 | |
| 176 | # Collect what was sent to the websocket |
| 177 | sent_payloads = [json.loads(c.args[0]) for c in mock_ws.send.call_args_list] |
| 178 | |
| 179 | # Initial send is the terminal size; everything after is stdin chars or our exit string. |
| 180 | # Verify "exit\n" was sent on EOF (one send per character) |
| 181 | assert {"data": "exit\n"} in sent_payloads, ( |
| 182 | f"Expected 'exit\\n' to be sent on piped EOF, got: {sent_payloads}" |
| 183 | ) |
| 184 | |
| 185 | # Verify we closed the websocket from the client side |
| 186 | mock_ws.close.assert_called_once() |
| 187 | |
| 188 | |
| 189 | @patch("colab_cli.console.os.get_terminal_size") |
| 190 | @patch("colab_cli.console.sys.stdin.isatty") |
| 191 | @patch("colab_cli.console.sys.stdin") |
| 192 | def test_read_stdin_eof_tty_does_not_close_ws( |
| 193 | mock_stdin, mock_isatty, mock_get_term_size |
| 194 | ): |
| 195 | """When stdin is a real TTY and read() returns empty (which happens on |
| 196 | Ctrl-D in raw mode), we should NOT inject 'exit\\n' or close the websocket |
| 197 | \u2014 the user is in interactive mode and may have intended Ctrl-D as a literal |
| 198 | char. The websocket lifecycle is owned by the remote shell in this case. |
| 199 | """ |
| 200 | import colab_cli.console as console_mod |
| 201 | |
| 202 | mock_isatty.return_value = True |
| 203 | # TTY EOF is rare but possible; should be passed through transparently |
| 204 | mock_stdin.read.side_effect = [""] |
| 205 | mock_get_term_size.return_value = os.terminal_size((80, 24)) |
| 206 | |
| 207 | mock_ws = MagicMock() |
| 208 | |
| 209 | class SyncThread: |
| 210 | def __init__(self, target, daemon=None): |
| 211 | self.target = target |
| 212 | |
| 213 | def start(self): |
| 214 | self.target() |
| 215 | |
| 216 | console_mod._is_running = True |
| 217 | with patch("colab_cli.console.threading.Thread", SyncThread): |
| 218 | on_open(mock_ws) |
| 219 | |
| 220 | sent_payloads = [json.loads(c.args[0]) for c in mock_ws.send.call_args_list] |
| 221 | assert {"data": "exit\n"} not in sent_payloads |
| 222 | mock_ws.close.assert_not_called() |