| 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 | |
| 22 | runner = CliRunner() |
| 23 | |
| 24 | |
| 25 | @pytest.fixture |
| 26 | def mock_store(mock_common_state): |
| 27 | return mock_common_state.store |
| 28 | |
| 29 | |
| 30 | @pytest.fixture |
| 31 | def mock_runtime_class(mocker): |
| 32 | # Patch it in the command module where it's used |
| 33 | return mocker.patch("colab_cli.commands.execution.ColabRuntime") |
| 34 | |
| 35 | |
| 36 | def test_cli_exec_file(mock_store, mock_runtime_class, mock_common_state, tmp_path): |
| 37 | mock_session = MagicMock() |
| 38 | mock_session.url = "http://url" |
| 39 | mock_session.token = "token123" |
| 40 | mock_session.name = "s1" |
| 41 | mock_session.kernel_id = None |
| 42 | mock_session.session_id = None |
| 43 | mock_store.get.return_value = mock_session |
| 44 | |
| 45 | mock_common_state.resolve_session.return_value = "s1" |
| 46 | mock_runtime = mock_runtime_class.return_value |
| 47 | mock_runtime.execute_code.return_value = [{"text": "hello\n"}] |
| 48 | |
| 49 | script = tmp_path / "script.py" |
| 50 | script.write_text("print('hello')") |
| 51 | |
| 52 | result = runner.invoke(app, ["exec", "-s", "s1", "-f", str(script)]) |
| 53 | assert result.exit_code == 0 |
| 54 | assert mock_session.last_execution[0] == str(script) |
| 55 | assert mock_session.last_execution[1] is None |
| 56 | assert mock_session.last_execution[2] is not None |
| 57 | mock_store.add.assert_called_with(mock_session) |
| 58 | mock_runtime.execute_code.assert_any_call( |
| 59 | "import os; os.makedirs('/content', exist_ok=True); os.chdir('/content')" |
| 60 | ) |
| 61 | mock_runtime.execute_code.assert_any_call( |
| 62 | "print('hello')", output_hook=ANY, timeout=30.0 |
| 63 | ) |
| 64 | |
| 65 | |
| 66 | def test_cli_exec_stdin(mock_store, mock_runtime_class, mock_common_state): |
| 67 | mock_session = MagicMock() |
| 68 | mock_session.name = "s1" |
| 69 | mock_session.url = "http://url" |
| 70 | mock_session.token = "token" |
| 71 | mock_session.kernel_id = None |
| 72 | mock_session.session_id = None |
| 73 | mock_store.get.return_value = mock_session |
| 74 | |
| 75 | mock_common_state.resolve_session.return_value = "s1" |
| 76 | mock_runtime = mock_runtime_class.return_value |
| 77 | mock_runtime.execute_code.return_value = [{"data": {"text/plain": "42"}}] |
| 78 | |
| 79 | result = runner.invoke(app, ["exec", "-s", "s1"], input="print(42)") |
| 80 | assert result.exit_code == 0 |
| 81 | assert mock_session.last_execution[0] == "stdin" |
| 82 | assert mock_session.last_execution[1] is None |
| 83 | assert mock_session.last_execution[2] is not None |
| 84 | mock_store.add.assert_called_with(mock_session) |
| 85 | mock_runtime.execute_code.assert_any_call( |
| 86 | "print(42)", output_hook=ANY, timeout=30.0 |
| 87 | ) |
| 88 | |
| 89 | |
| 90 | def test_cli_exec_env_injects_prelude( |
| 91 | mock_store, mock_runtime_class, mock_common_state |
| 92 | ): |
| 93 | mock_session = MagicMock() |
| 94 | mock_session.name = "s1" |
| 95 | mock_session.url = "http://url" |
| 96 | mock_session.token = "token" |
| 97 | mock_session.kernel_id = None |
| 98 | mock_session.session_id = None |
| 99 | mock_store.get.return_value = mock_session |
| 100 | |
| 101 | mock_common_state.resolve_session.return_value = "s1" |
| 102 | mock_runtime = mock_runtime_class.return_value |
| 103 | mock_runtime.execute_code.return_value = [] |
| 104 | |
| 105 | code = "import os\nprint(os.environ.get('HF_TOKEN'))" |
| 106 | result = runner.invoke( |
| 107 | app, ["exec", "-s", "s1", "--env", "HF_TOKEN=abc"], input=code |
| 108 | ) |
| 109 | |
| 110 | assert result.exit_code == 0, result.output |
| 111 | expected = "import os\nos.environ['HF_TOKEN'] = 'abc'\n" + code |
| 112 | mock_runtime.execute_code.assert_any_call(expected, output_hook=ANY, timeout=30.0) |
| 113 | |
| 114 | |
| 115 | def test_cli_exec_env_flags_accumulate_and_split_on_first_equals( |
| 116 | mock_store, mock_runtime_class, mock_common_state |
| 117 | ): |
| 118 | mock_session = MagicMock() |
| 119 | mock_session.name = "s1" |
| 120 | mock_session.url = "http://url" |
| 121 | mock_session.token = "token" |
| 122 | mock_session.kernel_id = None |
| 123 | mock_session.session_id = None |
| 124 | mock_store.get.return_value = mock_session |
| 125 | |
| 126 | mock_common_state.resolve_session.return_value = "s1" |
| 127 | mock_runtime = mock_runtime_class.return_value |
| 128 | mock_runtime.execute_code.return_value = [] |
| 129 | |
| 130 | result = runner.invoke( |
| 131 | app, |
| 132 | ["exec", "-s", "s1", "--env", "HF_TOKEN=abc", "--env", "B64=a=b=c"], |
| 133 | input="print('ok')", |
| 134 | ) |
| 135 | |
| 136 | assert result.exit_code == 0, result.output |
| 137 | expected = ( |
| 138 | "import os\n" |
| 139 | "os.environ['HF_TOKEN'] = 'abc'\n" |
| 140 | "os.environ['B64'] = 'a=b=c'\n" |
| 141 | "print('ok')" |
| 142 | ) |
| 143 | mock_runtime.execute_code.assert_any_call(expected, output_hook=ANY, timeout=30.0) |
| 144 | |
| 145 | |
| 146 | def test_cli_exec_env_escapes_tricky_literals( |
| 147 | mock_store, mock_runtime_class, mock_common_state |
| 148 | ): |
| 149 | mock_session = MagicMock() |
| 150 | mock_session.name = "s1" |
| 151 | mock_session.url = "http://url" |
| 152 | mock_session.token = "token" |
| 153 | mock_session.kernel_id = None |
| 154 | mock_session.session_id = None |
| 155 | mock_store.get.return_value = mock_session |
| 156 | |
| 157 | mock_common_state.resolve_session.return_value = "s1" |
| 158 | mock_runtime = mock_runtime_class.return_value |
| 159 | mock_runtime.execute_code.return_value = [] |
| 160 | |
| 161 | value = "quote'back\\slash=µ" |
| 162 | result = runner.invoke( |
| 163 | app, ["exec", "-s", "s1", "--env", f"TRICKY={value}"], input="print('ok')" |
| 164 | ) |
| 165 | |
| 166 | assert result.exit_code == 0, result.output |
| 167 | expected = f"import os\nos.environ['TRICKY'] = {value!r}\nprint('ok')" |
| 168 | mock_runtime.execute_code.assert_any_call(expected, output_hook=ANY, timeout=30.0) |
| 169 | |
| 170 | |
| 171 | def test_cli_exec_malformed_env_errors_before_session_resolution( |
| 172 | mock_runtime_class, mock_common_state |
| 173 | ): |
| 174 | result = runner.invoke( |
| 175 | app, ["exec", "-s", "s1", "--env", "HF_TOKEN"], input="print('ok')" |
| 176 | ) |
| 177 | |
| 178 | assert result.exit_code != 0 |
| 179 | assert "Expected KEY=VALUE" in result.output |
| 180 | mock_common_state.resolve_session.assert_not_called() |
| 181 | mock_runtime_class.assert_not_called() |
| 182 | |
| 183 | |
| 184 | def test_cli_exec_not_found(mock_common_state): |
| 185 | # Case where resolve_session fails |
| 186 | mock_common_state.resolve_session.side_effect = SystemExit(1) |
| 187 | result = runner.invoke(app, ["exec", "-s", "missing"]) |
| 188 | assert result.exit_code == 1 |
| 189 | |
| 190 | |
| 191 | def test_cli_exec_no_input(mock_store, mock_common_state, mocker): |
| 192 | mock_session = MagicMock() |
| 193 | mock_session.name = "s1" |
| 194 | mock_store.get.return_value = mock_session |
| 195 | |
| 196 | mock_common_state.resolve_session.return_value = "s1" |
| 197 | |
| 198 | # Mock is_stdin_tty to True to trigger the "No input provided" error |
| 199 | mocker.patch("colab_cli.commands.execution.is_stdin_tty", return_value=True) |
| 200 | |
| 201 | result = runner.invoke(app, ["exec", "-s", "s1"]) |
| 202 | assert result.exit_code == 1 |
| 203 | assert "No input provided" in result.output |
| 204 | |
| 205 | |
| 206 | @patch("colab_cli.commands.execution.handle_image") |
| 207 | def test_cli_exec_outputs( |
| 208 | mock_handle_image, mock_store, mock_runtime_class, mock_common_state |
| 209 | ): |
| 210 | mock_session = MagicMock() |
| 211 | mock_session.name = "s1" |
| 212 | mock_session.url = "http://url" |
| 213 | mock_session.token = "token" |
| 214 | mock_session.kernel_id = None |
| 215 | mock_session.session_id = None |
| 216 | mock_store.get.return_value = mock_session |
| 217 | |
| 218 | mock_common_state.resolve_session.return_value = "s1" |
| 219 | mock_runtime = mock_runtime_class.return_value |
| 220 | |
| 221 | # We need to simulate the output_hook being called because the command now relies on it |
| 222 | # for immediate output, although it also returns the final list. |
| 223 | def mock_execute_code(code, output_hook=None, **kwargs): |
| 224 | outputs = [ |
| 225 | {"data": {"image/png": "png_data"}}, |
| 226 | {"data": {"image/jpeg": "jpeg_data"}}, |
| 227 | {"output_type": "error", "ename": "ValueError", "evalue": "bad"}, |
| 228 | {"output_type": "error", "traceback": ["line1\n", "line2\n"]}, |
| 229 | ] |
| 230 | if output_hook: |
| 231 | for o in outputs: |
| 232 | output_hook(o) |
| 233 | return outputs |
| 234 | |
| 235 | mock_runtime.execute_code.side_effect = mock_execute_code |
| 236 | |
| 237 | result = runner.invoke(app, ["exec", "-s", "s1"], input="do_stuff()") |
| 238 | assert result.exit_code == 0 |
| 239 | |
| 240 | mock_handle_image.assert_any_call("png_data", "image/png", target_path=None) |
| 241 | mock_handle_image.assert_any_call("jpeg_data", "image/jpeg", target_path=None) |
| 242 | |
| 243 | assert "ValueError: bad\n" in result.stderr |
| 244 | assert "line1\nline2\n" in result.stderr |
| 245 | |
| 246 | |
| 247 | def test_cli_exec_empty_code(mock_runtime_class, mock_store, mock_common_state): |
| 248 | mock_session = MagicMock() |
| 249 | mock_session.name = "s1" |
| 250 | mock_session.url = "http://url" |
| 251 | mock_session.token = "token" |
| 252 | mock_session.kernel_id = None |
| 253 | mock_session.session_id = None |
| 254 | mock_store.get.return_value = mock_session |
| 255 | |
| 256 | mock_common_state.resolve_session.return_value = "s1" |
| 257 | result = runner.invoke(app, ["exec", "-s", "s1"], input=" \n ") |
| 258 | assert result.exit_code == 0 |
| 259 | |
| 260 | |
| 261 | def test_cli_exec_lost_session_prunes( |
| 262 | mock_runtime_class, mock_store, mock_common_state |
| 263 | ): |
| 264 | mock_session = MagicMock() |
| 265 | mock_session.name = "lost-sess" |
| 266 | mock_store.get.return_value = mock_session |
| 267 | mock_common_state.resolve_session.return_value = "lost-sess" |
| 268 | |
| 269 | mock_runtime = mock_runtime_class.return_value |
| 270 | # Simulate 404 during initialization |
| 271 | mock_runtime.execute_code.side_effect = Exception("404 Not Found") |
| 272 | |
| 273 | result = runner.invoke(app, ["exec", "-s", "lost-sess"], input="print(1)") |
| 274 | assert result.exit_code == 1 |
| 275 | assert "appears to be lost" in result.output |
| 276 | mock_common_state.prune_session.assert_called_once_with("lost-sess") |
| 277 | |
| 278 | |
| 279 | def test_cli_exec_timeout(mock_store, mock_runtime_class, mock_common_state, tmp_path): |
| 280 | mock_session = MagicMock() |
| 281 | mock_session.url = "http://url" |
| 282 | mock_session.token = "token123" |
| 283 | mock_session.name = "s1" |
| 284 | mock_session.kernel_id = None |
| 285 | mock_session.session_id = None |
| 286 | mock_store.get.return_value = mock_session |
| 287 | |
| 288 | mock_common_state.resolve_session.return_value = "s1" |
| 289 | mock_runtime = mock_runtime_class.return_value |
| 290 | mock_runtime.execute_code.return_value = [{"text": "hello\n"}] |
| 291 | |
| 292 | script = tmp_path / "script.py" |
| 293 | script.write_text("print('hello')") |
| 294 | |
| 295 | result = runner.invoke( |
| 296 | app, ["exec", "-s", "s1", "-f", str(script), "--timeout", "3600"] |
| 297 | ) |
| 298 | assert result.exit_code == 0 |
| 299 | mock_runtime.execute_code.assert_any_call( |
| 300 | "print('hello')", output_hook=ANY, timeout=3600.0 |
| 301 | ) |