| 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 | """Tests for `colab ssh` auto-create + --proxy-mode flag behavior. |
| 16 | |
| 17 | Bare `colab ssh` (no -s): ssh into the single existing session, auto-create one |
| 18 | when there are none, or error when there are several. `--gpu/--tpu` pass through |
| 19 | to an auto-created runtime (and are ignored otherwise); `--rm` stops a runtime |
| 20 | this command created. In --proxy-mode every flag still applies: `-s NAME` |
| 21 | creates the session if missing, and --rm stops the bridged session on |
| 22 | disconnect (installing SIGHUP/SIGTERM/SIGINT handlers so teardown survives the |
| 23 | way OpenSSH ends a ProxyCommand). |
| 24 | """ |
| 25 | |
| 26 | from unittest.mock import MagicMock |
| 27 | |
| 28 | from colab_cli.cli import app |
| 29 | from colab_cli.commands import ssh as ssh_module |
| 30 | import pytest |
| 31 | from typer.testing import CliRunner |
| 32 | import typer |
| 33 | |
| 34 | runner = CliRunner() |
| 35 | |
| 36 | |
| 37 | def _make_session( |
| 38 | name: str = "auto1", |
| 39 | url: str = "https://abc.colab.googleusercontent.com", |
| 40 | token: str = "TOK", |
| 41 | endpoint: str = "ep1", |
| 42 | ): |
| 43 | s = MagicMock() |
| 44 | s.name = name |
| 45 | s.url = url |
| 46 | s.token = token |
| 47 | s.endpoint = endpoint |
| 48 | return s |
| 49 | |
| 50 | |
| 51 | def _patch_interactive(mocker): |
| 52 | mocker.patch.object( |
| 53 | ssh_module, "_resolve_pubkey", return_value="ssh-ed25519 AAAA u@h" |
| 54 | ) |
| 55 | return mocker.patch.object( |
| 56 | ssh_module, "_run_interactive_ssh", return_value=0 |
| 57 | ) |
| 58 | |
| 59 | |
| 60 | def _patch_proxy(mocker): |
| 61 | mocker.patch.object( |
| 62 | ssh_module, "_resolve_pubkey", return_value="ssh-ed25519 AAAA u@h" |
| 63 | ) |
| 64 | mocker.patch.object( |
| 65 | ssh_module, "_connect_websocket", return_value=MagicMock() |
| 66 | ) |
| 67 | mocker.patch.object(ssh_module, "_bridge_proxy_mode", return_value=0) |
| 68 | |
| 69 | |
| 70 | # --- bare `colab ssh`: create / reuse / ambiguous --------------------------- |
| 71 | |
| 72 | |
| 73 | @pytest.mark.parametrize( |
| 74 | ("sessions", "resolve_raises", "gpu", "expect_create", "expect_ok"), |
| 75 | [ |
| 76 | ({}, False, None, True, True), |
| 77 | ({}, False, "T4", True, True), |
| 78 | ({"only": 1}, False, None, False, True), |
| 79 | ({"a": 1, "b": 1}, True, None, False, False), |
| 80 | ], |
| 81 | ids=["zero-creates", "zero-creates-gpu", "one-reuses", "many-errors"], |
| 82 | ) |
| 83 | def test_bare_ssh_session_resolution( |
| 84 | mock_common_state, |
| 85 | mocker, |
| 86 | sessions, |
| 87 | resolve_raises, |
| 88 | gpu, |
| 89 | expect_create, |
| 90 | expect_ok, |
| 91 | ): |
| 92 | mock_common_state.store.list.return_value = sessions |
| 93 | sess = _make_session() |
| 94 | mock_common_state.store.get.return_value = sess |
| 95 | if resolve_raises: |
| 96 | mock_common_state.resolve_session.side_effect = typer.Exit(1) |
| 97 | else: |
| 98 | mock_common_state.resolve_session.return_value = "only" |
| 99 | |
| 100 | new = mocker.patch("colab_cli.commands.session.new") |
| 101 | interactive = _patch_interactive(mocker) |
| 102 | |
| 103 | args = ["ssh"] + (["--gpu", gpu] if gpu else []) |
| 104 | result = runner.invoke(app, args) |
| 105 | |
| 106 | if expect_ok: |
| 107 | assert result.exit_code == 0 |
| 108 | interactive.assert_called_once_with(sess, None) |
| 109 | else: |
| 110 | assert result.exit_code != 0 |
| 111 | interactive.assert_not_called() |
| 112 | |
| 113 | if expect_create: |
| 114 | new.assert_called_once() |
| 115 | assert new.call_args.kwargs.get("gpu") == gpu |
| 116 | assert new.call_args.kwargs.get("tpu") is None |
| 117 | assert new.call_args.kwargs.get("high_mem") is False |
| 118 | else: |
| 119 | new.assert_not_called() |
| 120 | |
| 121 | |
| 122 | def test_bare_ssh_forwards_high_mem_on_autocreate(mock_common_state, mocker): |
| 123 | mock_common_state.store.list.return_value = {} |
| 124 | sess = _make_session() |
| 125 | mock_common_state.store.get.return_value = sess |
| 126 | new = mocker.patch("colab_cli.commands.session.new") |
| 127 | _patch_interactive(mocker) |
| 128 | |
| 129 | result = runner.invoke(app, ["ssh", "--gpu", "A100", "--high-mem"]) |
| 130 | assert result.exit_code == 0 |
| 131 | new.assert_called_once() |
| 132 | assert new.call_args.kwargs == { |
| 133 | "session": new.call_args.kwargs["session"], |
| 134 | "gpu": "A100", |
| 135 | "tpu": None, |
| 136 | "high_mem": True, |
| 137 | } |
| 138 | |
| 139 | |
| 140 | # --- --proxy-mode -s NAME: create-if-missing / reuse ------------------------ |
| 141 | |
| 142 | |
| 143 | @pytest.mark.parametrize( |
| 144 | ("exists", "gpu", "expect_new"), |
| 145 | [ |
| 146 | (False, None, True), |
| 147 | (False, "T4", True), |
| 148 | (True, None, False), |
| 149 | ], |
| 150 | ids=["missing-creates", "missing-creates-gpu", "existing-reuses"], |
| 151 | ) |
| 152 | def test_proxy_mode_create_or_reuse( |
| 153 | mock_common_state, mocker, exists, gpu, expect_new |
| 154 | ): |
| 155 | created = _make_session("colab") |
| 156 | if exists: |
| 157 | mock_common_state.store.get.return_value = created |
| 158 | mock_common_state.resolve_session.return_value = "colab" |
| 159 | else: |
| 160 | mock_common_state.store.get.return_value = None |
| 161 | |
| 162 | def after_new(*a, **k): |
| 163 | mock_common_state.store.get.return_value = created |
| 164 | |
| 165 | new = mocker.patch("colab_cli.commands.session.new", side_effect=after_new) |
| 166 | _patch_proxy(mocker) |
| 167 | |
| 168 | args = ["ssh", "--proxy-mode", "-s", "colab"] |
| 169 | if gpu: |
| 170 | args += ["--gpu", gpu] |
| 171 | result = runner.invoke(app, args) |
| 172 | assert result.exit_code == 0 |
| 173 | |
| 174 | if expect_new: |
| 175 | new.assert_called_once() |
| 176 | assert new.call_args.kwargs.get("session") == "colab" |
| 177 | assert new.call_args.kwargs.get("gpu") == gpu |
| 178 | else: |
| 179 | new.assert_not_called() |
| 180 | |
| 181 | |
| 182 | def test_proxy_mode_no_session_does_not_autocreate(mock_common_state, mocker): |
| 183 | """--proxy-mode with no -s does not auto-create (nothing to name).""" |
| 184 | mock_common_state.store.list.return_value = {} |
| 185 | mock_common_state.resolve_session.side_effect = typer.Exit(2) |
| 186 | new = mocker.patch("colab_cli.commands.session.new") |
| 187 | connect = mocker.patch.object(ssh_module, "_connect_websocket") |
| 188 | mocker.patch.object( |
| 189 | ssh_module, "_resolve_pubkey", return_value="ssh-ed25519 AAAA u@h" |
| 190 | ) |
| 191 | |
| 192 | result = runner.invoke(app, ["ssh", "--proxy-mode"]) |
| 193 | assert result.exit_code != 0 |
| 194 | new.assert_not_called() |
| 195 | connect.assert_not_called() |
| 196 | |
| 197 | |
| 198 | # --- --gpu/--tpu ignored when not creating ---------------------------------- |
| 199 | |
| 200 | |
| 201 | @pytest.mark.parametrize( |
| 202 | ("extra_args", "expect_msg"), |
| 203 | [ |
| 204 | ( |
| 205 | ["--proxy-mode", "-s", "colab", "--gpu", "T4"], |
| 206 | "only applies to a created runtime", |
| 207 | ), |
| 208 | ( |
| 209 | ["-s", "colab", "--tpu", "v5e1"], |
| 210 | "only applies to a created runtime", |
| 211 | ), |
| 212 | ], |
| 213 | ids=["proxy-existing", "interactive-reuse"], |
| 214 | ) |
| 215 | def test_gpu_tpu_ignored_when_not_creating( |
| 216 | mock_common_state, mocker, extra_args, expect_msg |
| 217 | ): |
| 218 | sess = _make_session("colab") |
| 219 | mock_common_state.store.get.return_value = sess |
| 220 | mock_common_state.store.list.return_value = {"colab": sess} |
| 221 | mock_common_state.resolve_session.return_value = "colab" |
| 222 | _patch_proxy(mocker) |
| 223 | mocker.patch.object(ssh_module, "_run_interactive_ssh", return_value=0) |
| 224 | |
| 225 | result = runner.invoke(app, ["ssh", *extra_args]) |
| 226 | assert result.exit_code == 0 |
| 227 | assert expect_msg in result.stderr |
| 228 | |
| 229 | |
| 230 | # --- --rm teardown ---------------------------------------------------------- |
| 231 | |
| 232 | |
| 233 | @pytest.mark.parametrize( |
| 234 | ("rm", "expect_stop"), |
| 235 | [(True, True), (False, False)], |
| 236 | ids=["rm-stops", "no-rm-keeps"], |
| 237 | ) |
| 238 | def test_proxy_mode_rm_teardown(mock_common_state, mocker, rm, expect_stop): |
| 239 | created = _make_session("colab-ephem") |
| 240 | mock_common_state.store.get.return_value = None |
| 241 | |
| 242 | def after_new(*a, **k): |
| 243 | mock_common_state.store.get.return_value = created |
| 244 | |
| 245 | mocker.patch("colab_cli.commands.session.new", side_effect=after_new) |
| 246 | stop = mocker.patch("colab_cli.commands.session.stop") |
| 247 | mocker.patch("signal.signal") # don't install real handlers during tests |
| 248 | _patch_proxy(mocker) |
| 249 | |
| 250 | args = ["ssh", "--proxy-mode", "-s", "colab-ephem"] |
| 251 | if rm: |
| 252 | args.append("--rm") |
| 253 | result = runner.invoke(app, args) |
| 254 | assert result.exit_code == 0 |
| 255 | if expect_stop: |
| 256 | stop.assert_called_once_with(session="colab-ephem") |
| 257 | else: |
| 258 | stop.assert_not_called() |
| 259 | |
| 260 | |
| 261 | @pytest.mark.parametrize( |
| 262 | ("has_existing", "expect_stop"), |
| 263 | [(False, True), (True, False)], |
| 264 | ids=["autocreated-stops", "reused-keeps"], |
| 265 | ) |
| 266 | def test_interactive_rm_teardown( |
| 267 | mock_common_state, mocker, has_existing, expect_stop |
| 268 | ): |
| 269 | """Interactive --rm stops only a runtime `colab ssh` auto-created.""" |
| 270 | if has_existing: |
| 271 | sess = _make_session("only") |
| 272 | mock_common_state.store.list.return_value = {"only": sess} |
| 273 | mock_common_state.resolve_session.return_value = "only" |
| 274 | else: |
| 275 | sess = _make_session("auto-rm") |
| 276 | mock_common_state.store.list.return_value = {} |
| 277 | mock_common_state.store.get.return_value = sess |
| 278 | |
| 279 | mocker.patch("colab_cli.commands.session.new") |
| 280 | stop = mocker.patch("colab_cli.commands.session.stop") |
| 281 | _patch_interactive(mocker) |
| 282 | |
| 283 | result = runner.invoke(app, ["ssh", "--rm"]) |
| 284 | assert result.exit_code == 0 |
| 285 | if expect_stop: |
| 286 | stop.assert_called_once_with(session="auto-rm") |
| 287 | else: |
| 288 | stop.assert_not_called() |
| 289 | |
| 290 | |
| 291 | # --- signal-handler installation (proxy-mode --rm) -------------------------- |
| 292 | |
| 293 | |
| 294 | @pytest.mark.parametrize( |
| 295 | ("rm", "expect_installed"), |
| 296 | [(True, True), (False, False)], |
| 297 | ids=["rm-installs", "no-rm-none"], |
| 298 | ) |
| 299 | def test_proxy_mode_signal_handler_installation( |
| 300 | mock_common_state, mocker, rm, expect_installed |
| 301 | ): |
| 302 | """--rm installs SIGHUP/SIGTERM/SIGINT handlers so teardown runs when |
| 303 | OpenSSH HUPs the ProxyCommand on disconnect; without --rm, none.""" |
| 304 | import signal as _signal |
| 305 | |
| 306 | mock_common_state.store.get.return_value = _make_session("colab") |
| 307 | mock_common_state.resolve_session.return_value = "colab" |
| 308 | sigmock = mocker.patch("signal.signal") |
| 309 | mocker.patch("colab_cli.commands.session.stop") |
| 310 | _patch_proxy(mocker) |
| 311 | |
| 312 | args = ["ssh", "--proxy-mode", "-s", "colab"] |
| 313 | if rm: |
| 314 | args.append("--rm") |
| 315 | result = runner.invoke(app, args) |
| 316 | assert result.exit_code == 0 |
| 317 | if expect_installed: |
| 318 | registered = {c.args[0] for c in sigmock.call_args_list} |
| 319 | assert {_signal.SIGHUP, _signal.SIGTERM, _signal.SIGINT} <= registered |
| 320 | else: |
| 321 | sigmock.assert_not_called() |
| 322 | |
| 323 | |
| 324 | # --- help -------------------------------------------------------------------- |
| 325 | |
| 326 | |
| 327 | def test_ssh_help_advertises_autocreate_flags(): |
| 328 | """`colab ssh --help` advertises --rm, --gpu, and --tpu.""" |
| 329 | result = runner.invoke(app, ["ssh", "--help"]) |
| 330 | assert result.exit_code == 0 |
| 331 | for flag in ("--rm", "--gpu", "--tpu"): |
| 332 | assert flag in result.output |