| 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 | from unittest.mock import MagicMock, patch, ANY |
| 16 | |
| 17 | import pytest |
| 18 | from typer.testing import CliRunner |
| 19 | |
| 20 | from colab_cli.cli import app |
| 21 | from colab_cli.repl import ColabREPL |
| 22 | |
| 23 | runner = CliRunner() |
| 24 | |
| 25 | |
| 26 | @pytest.fixture |
| 27 | def mock_store(mock_common_state): |
| 28 | return mock_common_state.store |
| 29 | |
| 30 | |
| 31 | @pytest.fixture |
| 32 | def mock_runtime_class(mocker): |
| 33 | # Patch it in the command module where it's used |
| 34 | return mocker.patch("colab_cli.commands.execution.ColabRuntime") |
| 35 | |
| 36 | |
| 37 | @patch("colab_cli.repl.handle_image") |
| 38 | def test_repl_display_output(mock_handle_image, capsys): |
| 39 | runtime = MagicMock() |
| 40 | repl_inst = ColabREPL(runtime) |
| 41 | |
| 42 | outputs = [ |
| 43 | {"text": "hello"}, |
| 44 | {"data": {"image/png": "png_data", "text/plain": "<Figure size>"}}, |
| 45 | {"data": {"image/jpeg": "jpeg_data", "text/plain": "other_text"}}, |
| 46 | {"output_type": "error", "ename": "ValueError", "evalue": "bad"}, |
| 47 | {"output_type": "error", "traceback": ["line1\n", "line2\n"]}, |
| 48 | ] |
| 49 | |
| 50 | for o in outputs: |
| 51 | repl_inst.display_output(o) |
| 52 | |
| 53 | mock_handle_image.assert_any_call("png_data", "image/png", target_path=None) |
| 54 | mock_handle_image.assert_any_call("jpeg_data", "image/jpeg", target_path=None) |
| 55 | |
| 56 | captured = capsys.readouterr() |
| 57 | assert "hello" in captured.out |
| 58 | assert "other_text" in captured.out |
| 59 | assert "ValueError: bad" in captured.out |
| 60 | assert "line1\nline2\n" in captured.out |
| 61 | |
| 62 | |
| 63 | def test_repl_execute(mock_store, mock_common_state): |
| 64 | runtime = MagicMock() |
| 65 | |
| 66 | mock_session = MagicMock() |
| 67 | mock_store.get.return_value = mock_session |
| 68 | |
| 69 | def mock_execute_code(code, output_hook=None, **kwargs): |
| 70 | o = {"text": "done"} |
| 71 | if output_hook: |
| 72 | output_hook(o) |
| 73 | return [o] |
| 74 | |
| 75 | runtime.execute_code.side_effect = mock_execute_code |
| 76 | repl_inst = ColabREPL(runtime, session_name="s1") |
| 77 | |
| 78 | with patch.object(repl_inst, "display_output") as mock_display: |
| 79 | repl_inst.execute("print(1)") |
| 80 | mock_display.assert_called_once_with({"text": "done"}) |
| 81 | assert repl_inst.repl_history[0]["input"] == "print(1)" |
| 82 | |
| 83 | assert mock_session.last_execution[0] == "REPL" |
| 84 | assert mock_session.last_execution[1] is None |
| 85 | assert mock_session.last_execution[2] is not None |
| 86 | mock_store.add.assert_called_with(mock_session) |
| 87 | |
| 88 | |
| 89 | def test_cli_repl_interactive( |
| 90 | mock_runtime_class, mock_store, mock_common_state, mocker |
| 91 | ): |
| 92 | mock_session = MagicMock() |
| 93 | mock_session.name = "s1" |
| 94 | mock_session.url = "http://url" |
| 95 | mock_session.token = "token" |
| 96 | mock_session.kernel_id = None |
| 97 | mock_session.session_id = None |
| 98 | mock_store.get.return_value = mock_session |
| 99 | |
| 100 | mock_common_state.resolve_session.return_value = "s1" |
| 101 | |
| 102 | # Mock is_stdin_tty to True to follow the interactive path |
| 103 | mocker.patch("colab_cli.commands.execution.is_stdin_tty", return_value=True) |
| 104 | |
| 105 | # Simulate TTY for interactive REPL |
| 106 | with patch("colab_cli.repl.ColabREPL") as mock_repl_class: |
| 107 | # Mock run to prevent infinite loop or errors |
| 108 | mock_repl_class.return_value.run.return_value = None |
| 109 | runner.invoke(app, ["repl", "-s", "s1"]) |
| 110 | assert mock_repl_class.called |
| 111 | |
| 112 | |
| 113 | def test_cli_repl_piped(mock_runtime_class, mock_store, mock_common_state): |
| 114 | mock_session = MagicMock() |
| 115 | mock_session.name = "s1" |
| 116 | mock_session.url = "http://url" |
| 117 | mock_session.token = "token" |
| 118 | mock_session.kernel_id = None |
| 119 | mock_session.session_id = None |
| 120 | mock_store.get.return_value = mock_session |
| 121 | |
| 122 | mock_runtime = mock_runtime_class.return_value |
| 123 | mock_runtime.execute_code.return_value = [{"text": "done piped"}] |
| 124 | |
| 125 | mock_common_state.resolve_session.return_value = "s1" |
| 126 | result = runner.invoke(app, ["repl", "-s", "s1"], input="print(1)") |
| 127 | assert result.exit_code == 0 |
| 128 | assert mock_session.last_execution[0] == "stdin" |
| 129 | assert mock_session.last_execution[2] is not None |
| 130 | mock_store.add.assert_called_with(mock_session) |
| 131 | mock_runtime.execute_code.assert_any_call("print(1)", output_hook=ANY) |
| 132 | |
| 133 | |
| 134 | def test_cli_repl_missing_session(mock_common_state): |
| 135 | mock_common_state.resolve_session.side_effect = SystemExit(1) |
| 136 | result = runner.invoke(app, ["repl", "-s", "missing"]) |
| 137 | assert result.exit_code == 1 |
| 138 | |
| 139 | |
| 140 | def test_cli_repl_piped_empty(mock_runtime_class, mock_store, mock_common_state): |
| 141 | mock_session = MagicMock() |
| 142 | mock_session.name = "s1" |
| 143 | mock_session.url = "http://url" |
| 144 | mock_session.token = "token" |
| 145 | mock_session.kernel_id = None |
| 146 | mock_session.session_id = None |
| 147 | mock_store.get.return_value = mock_session |
| 148 | |
| 149 | mock_common_state.resolve_session.return_value = "s1" |
| 150 | result = runner.invoke(app, ["repl", "-s", "s1"], input=" \n ") |
| 151 | assert result.exit_code == 0 |
| 152 | |
| 153 | |
| 154 | def test_repl_print_info_error(capsys): |
| 155 | repl_inst = ColabREPL(MagicMock()) |
| 156 | repl_inst.print_info("info_msg") |
| 157 | repl_inst.print_error("err_msg") |
| 158 | |
| 159 | |
| 160 | @patch("colab_cli.repl.handle_image") |
| 161 | def test_repl_display_output_image_suppress_text(mock_handle_image, capsys): |
| 162 | repl_inst = ColabREPL(MagicMock()) |
| 163 | output = {"data": {"image/png": "png_data", "text/plain": "<Figure size>"}} |
| 164 | repl_inst.display_output(output) |
| 165 | mock_handle_image.assert_called_once_with("png_data", "image/png", target_path=None) |
| 166 | |
| 167 | captured = capsys.readouterr() |
| 168 | assert "<Figure size>" not in captured.out |
| 169 | |
| 170 | |
| 171 | def test_repl_execute_error(capsys): |
| 172 | runtime = MagicMock() |
| 173 | runtime.execute_code.side_effect = Exception("Kernel ded") |
| 174 | repl_inst = ColabREPL(runtime) |
| 175 | |
| 176 | repl_inst.execute("print(1)") |
| 177 | |
| 178 | captured = capsys.readouterr() |
| 179 | assert "Kernel ded" in captured.out |
| 180 | |
| 181 | |
| 182 | def test_repl_run_quit_aliases(mocker): |
| 183 | runtime = MagicMock() |
| 184 | repl_inst = ColabREPL(runtime) |
| 185 | repl_inst.session = MagicMock() |
| 186 | |
| 187 | # Test /quit |
| 188 | repl_inst.session.prompt.side_effect = ["/quit"] |
| 189 | repl_inst.run() |
| 190 | assert runtime.stop.called |
| 191 | |
| 192 | # Test quit() |
| 193 | runtime.stop.reset_mock() |
| 194 | repl_inst.session.prompt.side_effect = ["quit()"] |
| 195 | repl_inst.run() |
| 196 | assert runtime.stop.called |
| 197 | |
| 198 | # Test exit() |
| 199 | runtime.stop.reset_mock() |
| 200 | repl_inst.session.prompt.side_effect = ["exit()"] |
| 201 | repl_inst.run() |
| 202 | assert runtime.stop.called |
| 203 | |
| 204 | |
| 205 | def test_repl_run_misc_inputs(mocker, capsys): |
| 206 | runtime = MagicMock() |
| 207 | repl_inst = ColabREPL(runtime) |
| 208 | repl_inst.session = MagicMock() |
| 209 | |
| 210 | # None result, empty string, then exit() |
| 211 | repl_inst.session.prompt.side_effect = [None, " ", "exit()"] |
| 212 | repl_inst.run() |
| 213 | assert runtime.stop.called |
| 214 | |
| 215 | |
| 216 | def test_repl_run_exceptions(mocker, capsys): |
| 217 | runtime = MagicMock() |
| 218 | repl_inst = ColabREPL(runtime) |
| 219 | repl_inst.session = MagicMock() |
| 220 | |
| 221 | # EOFError |
| 222 | repl_inst.session.prompt.side_effect = [EOFError()] |
| 223 | repl_inst.run() |
| 224 | assert "Goodbye!" in capsys.readouterr().out |
| 225 | |
| 226 | # KeyboardInterrupt |
| 227 | repl_inst.session.prompt.side_effect = [KeyboardInterrupt(), "exit()"] |
| 228 | repl_inst.run() |
| 229 | |
| 230 | # Generic Exception |
| 231 | repl_inst.session.prompt.side_effect = [Exception("ouch"), "exit()"] |
| 232 | repl_inst.run() |
| 233 | assert "REPL Error: ouch" in capsys.readouterr().out |
| 234 | |
| 235 | |
| 236 | def test_repl_run_executes_code(mocker): |
| 237 | runtime = MagicMock() |
| 238 | repl_inst = ColabREPL(runtime) |
| 239 | repl_inst.session = MagicMock() |
| 240 | |
| 241 | with patch.object(repl_inst, "execute") as mock_execute: |
| 242 | repl_inst.session.prompt.side_effect = ["print(1)", "/quit"] |
| 243 | repl_inst.run() |
| 244 | mock_execute.assert_called_once_with("print(1)") |
| 245 | |
| 246 | |
| 247 | def test_repl_execute_with_history(mock_store): |
| 248 | runtime = MagicMock() |
| 249 | history_logger = MagicMock() |
| 250 | repl_inst = ColabREPL(runtime, session_name="s1", history_logger=history_logger) |
| 251 | |
| 252 | repl_inst.execute("print(1)") |
| 253 | assert history_logger.log_event.called |
| 254 | |
| 255 | |
| 256 | def test_repl_key_bindings(mocker): |
| 257 | runtime = MagicMock() |
| 258 | repl_inst = ColabREPL(runtime) |
| 259 | |
| 260 | # Trigger 'enter' binding (which is mapped to c-m in prompt_toolkit) |
| 261 | mock_event = MagicMock() |
| 262 | repl_inst.kb.get_bindings_for_keys(("c-m",))[0].handler(mock_event) |
| 263 | assert mock_event.current_buffer.validate_and_handle.called |
| 264 | |
| 265 | # Trigger 'c-j' binding |
| 266 | mock_event = MagicMock() |
| 267 | repl_inst.kb.get_bindings_for_keys(("c-j",))[0].handler(mock_event) |
| 268 | mock_event.current_buffer.insert_text.assert_called_with("\n") |