| 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 | """Lifecycle & dispatch guarantees for `colab ssh`. |
| 16 | |
| 17 | Complements the unit suite (test_ssh.py), the autocreate/flag suite |
| 18 | (test_ssh_autocreate.py), and the real-wire suite (test_ssh_wire_contract.py) |
| 19 | by pinning the command's *lifecycle* contracts -- the same class of guarantees |
| 20 | test_run.py enforces for `colab run`: |
| 21 | |
| 22 | * A. --rm teardown runs even when the bridge/shell raises (try/finally). |
| 23 | * B. the ssh/bridge exit code propagates to the process exit code. |
| 24 | * C. in --proxy-mode, create + --rm chatter stays on stderr so stdout remains |
| 25 | the clean ssh byte stream (a dropped redirect would corrupt every real |
| 26 | connection yet pass the mocked suite -- cf. test_ssh_wire_contract.py). |
| 27 | * D. auto-create failure aborts before any WebSocket connect (don't burn a VM |
| 28 | then fail). |
| 29 | * E. --gpu/--tpu are forwarded verbatim to `colab new` (which owns precedence). |
| 30 | * F. --rm teardown is idempotent across the signal path and the finally path. |
| 31 | * G. --proxy-mode --rm stops a *reused* session too, not just an auto-created |
| 32 | one. |
| 33 | """ |
| 34 | |
| 35 | from unittest.mock import MagicMock |
| 36 | |
| 37 | from colab_cli.cli import app |
| 38 | from colab_cli.commands import ssh as ssh_module |
| 39 | import pytest |
| 40 | import typer |
| 41 | from typer.testing import CliRunner |
| 42 | |
| 43 | runner = CliRunner() |
| 44 | |
| 45 | |
| 46 | def _make_session( |
| 47 | name: str = "s1", |
| 48 | url: str = "https://abc.colab.googleusercontent.com", |
| 49 | token: str = "TOK", |
| 50 | endpoint: str = "ep", |
| 51 | ): |
| 52 | s = MagicMock() |
| 53 | s.name = name |
| 54 | s.url = url |
| 55 | s.token = token |
| 56 | s.endpoint = endpoint |
| 57 | return s |
| 58 | |
| 59 | |
| 60 | def _patch_proxy(mocker): |
| 61 | """Stub the proxy-mode I/O seams (pubkey, connect, bridge).""" |
| 62 | mocker.patch.object(ssh_module, "_resolve_pubkey", return_value="pk") |
| 63 | mocker.patch.object( |
| 64 | ssh_module, "_connect_websocket", return_value=MagicMock() |
| 65 | ) |
| 66 | mocker.patch.object(ssh_module, "_bridge_proxy_mode", return_value=0) |
| 67 | |
| 68 | |
| 69 | # --- A. --rm teardown survives an exception (try/finally guarantee) ---------- |
| 70 | |
| 71 | |
| 72 | def test_proxy_mode_rm_stops_even_if_bridge_raises(mock_common_state, mocker): |
| 73 | """If _bridge_proxy_mode raises, the --rm stop still runs (finally).""" |
| 74 | sess = _make_session("colab-ephem") |
| 75 | mock_common_state.store.get.return_value = sess |
| 76 | mock_common_state.resolve_session.return_value = "colab-ephem" |
| 77 | mocker.patch("signal.signal") |
| 78 | stop = mocker.patch("colab_cli.commands.session.stop") |
| 79 | mocker.patch.object(ssh_module, "_resolve_pubkey", return_value="pk") |
| 80 | mocker.patch.object( |
| 81 | ssh_module, "_connect_websocket", return_value=MagicMock() |
| 82 | ) |
| 83 | mocker.patch.object( |
| 84 | ssh_module, "_bridge_proxy_mode", side_effect=RuntimeError("ws died") |
| 85 | ) |
| 86 | |
| 87 | result = runner.invoke( |
| 88 | app, ["ssh", "--proxy-mode", "-s", "colab-ephem", "--rm"] |
| 89 | ) |
| 90 | assert result.exit_code != 0 |
| 91 | stop.assert_called_once_with(session="colab-ephem") |
| 92 | |
| 93 | |
| 94 | def test_interactive_rm_stops_even_if_shell_raises(mock_common_state, mocker): |
| 95 | """If _run_interactive_ssh raises, an auto-created runtime is still |
| 96 | stopped (finally).""" |
| 97 | sess = _make_session("auto-rm") |
| 98 | mock_common_state.store.list.return_value = {} # empty -> auto-create |
| 99 | mock_common_state.store.get.return_value = sess |
| 100 | mocker.patch("colab_cli.commands.session.new") |
| 101 | stop = mocker.patch("colab_cli.commands.session.stop") |
| 102 | mocker.patch.object(ssh_module, "_resolve_pubkey", return_value="pk") |
| 103 | mocker.patch.object( |
| 104 | ssh_module, "_run_interactive_ssh", side_effect=RuntimeError("boom") |
| 105 | ) |
| 106 | |
| 107 | result = runner.invoke(app, ["ssh", "--rm"]) |
| 108 | assert result.exit_code != 0 |
| 109 | stop.assert_called_once_with(session="auto-rm") |
| 110 | |
| 111 | |
| 112 | # --- B. exit-code propagation ------------------------------------------------ |
| 113 | |
| 114 | |
| 115 | @pytest.mark.parametrize("code", [0, 1, 255], ids=["ok", "err", "ssh-fail"]) |
| 116 | def test_interactive_exit_code_propagates(mock_common_state, mocker, code): |
| 117 | """The interactive ssh subprocess's exit code becomes the CLI exit code.""" |
| 118 | sess = _make_session("s1") |
| 119 | mock_common_state.store.get.return_value = sess |
| 120 | mock_common_state.resolve_session.return_value = "s1" |
| 121 | mocker.patch.object(ssh_module, "_resolve_pubkey", return_value="pk") |
| 122 | mocker.patch.object(ssh_module, "_run_interactive_ssh", return_value=code) |
| 123 | |
| 124 | result = runner.invoke(app, ["ssh", "-s", "s1"]) |
| 125 | assert result.exit_code == code |
| 126 | |
| 127 | |
| 128 | @pytest.mark.parametrize("code", [0, 42], ids=["ok", "nonzero"]) |
| 129 | def test_proxy_mode_exit_code_propagates(mock_common_state, mocker, code): |
| 130 | """The proxy-mode bridge's return code becomes the CLI exit code.""" |
| 131 | sess = _make_session("s1") |
| 132 | mock_common_state.store.get.return_value = sess |
| 133 | mock_common_state.resolve_session.return_value = "s1" |
| 134 | _patch_proxy(mocker) |
| 135 | mocker.patch.object(ssh_module, "_bridge_proxy_mode", return_value=code) |
| 136 | |
| 137 | result = runner.invoke(app, ["ssh", "--proxy-mode", "-s", "s1"]) |
| 138 | assert result.exit_code == code |
| 139 | |
| 140 | |
| 141 | # --- C. --proxy-mode keeps stdout clean (byte-stream integrity) -------------- |
| 142 | |
| 143 | |
| 144 | def test_proxy_select_routes_create_output_to_stderr(mocker, capsys): |
| 145 | """Auto-create chatter must land on stderr, never stdout -- in |
| 146 | --proxy-mode stdout IS the ssh byte stream.""" |
| 147 | sess = _make_session("newname") |
| 148 | mocker.patch.object(ssh_module, "_session_exists", return_value=False) |
| 149 | mocker.patch.object(ssh_module, "_resolve_session", return_value=sess) |
| 150 | mocker.patch("colab_cli.commands.session.new") |
| 151 | |
| 152 | s, created = ssh_module._select_proxy_session("newname", "T4", None) |
| 153 | assert created is True and s is sess |
| 154 | |
| 155 | captured = capsys.readouterr() |
| 156 | assert "Creating runtime" in captured.err |
| 157 | assert "Creating runtime" not in captured.out |
| 158 | |
| 159 | |
| 160 | def test_proxy_bridge_routes_rm_output_to_stderr(mocker, capsys): |
| 161 | """--rm stop chatter must land on stderr, never stdout.""" |
| 162 | sess = _make_session("colab") |
| 163 | mocker.patch("signal.signal") |
| 164 | mocker.patch.object(ssh_module, "_resolve_pubkey", return_value="pk") |
| 165 | mocker.patch.object( |
| 166 | ssh_module, "_connect_websocket", return_value=MagicMock() |
| 167 | ) |
| 168 | mocker.patch.object(ssh_module, "_bridge_proxy_mode", return_value=0) |
| 169 | |
| 170 | def stop_echo(session=None): |
| 171 | typer.echo("STOP-MARKER") |
| 172 | |
| 173 | mocker.patch("colab_cli.commands.session.stop", side_effect=stop_echo) |
| 174 | |
| 175 | rc = ssh_module._run_proxy_bridge(sess, None, rm=True) |
| 176 | assert rc == 0 |
| 177 | |
| 178 | captured = capsys.readouterr() |
| 179 | assert "STOP-MARKER" in captured.err |
| 180 | assert "STOP-MARKER" not in captured.out |
| 181 | |
| 182 | |
| 183 | # --- D. auto-create failure aborts before any WebSocket connect -------------- |
| 184 | |
| 185 | |
| 186 | def test_proxy_mode_autocreate_failure_skips_connect(mock_common_state, mocker): |
| 187 | """A failed `colab new` in --proxy-mode must not proceed to connect.""" |
| 188 | mock_common_state.store.get.return_value = None # session missing |
| 189 | mocker.patch("colab_cli.commands.session.new", side_effect=typer.Exit(1)) |
| 190 | connect = mocker.patch.object(ssh_module, "_connect_websocket") |
| 191 | mocker.patch.object(ssh_module, "_resolve_pubkey", return_value="pk") |
| 192 | |
| 193 | result = runner.invoke(app, ["ssh", "--proxy-mode", "-s", "new"]) |
| 194 | assert result.exit_code != 0 |
| 195 | connect.assert_not_called() |
| 196 | |
| 197 | |
| 198 | def test_interactive_autocreate_failure_skips_ssh(mock_common_state, mocker): |
| 199 | """A failed `colab new` in interactive mode must not spawn ssh.""" |
| 200 | mock_common_state.store.list.return_value = {} # empty -> auto-create |
| 201 | mocker.patch("colab_cli.commands.session.new", side_effect=typer.Exit(1)) |
| 202 | interactive = mocker.patch.object(ssh_module, "_run_interactive_ssh") |
| 203 | mocker.patch.object(ssh_module, "_resolve_pubkey", return_value="pk") |
| 204 | |
| 205 | result = runner.invoke(app, ["ssh"]) |
| 206 | assert result.exit_code != 0 |
| 207 | interactive.assert_not_called() |
| 208 | |
| 209 | |
| 210 | # --- E. --gpu + --tpu are both forwarded to `colab new` ---------------------- |
| 211 | |
| 212 | |
| 213 | def test_gpu_and_tpu_both_forwarded_to_new(mock_common_state, mocker): |
| 214 | """`colab ssh --gpu T4 --tpu v5e1` forwards both to `colab new`, which |
| 215 | resolves precedence -- SSH does not silently drop either.""" |
| 216 | sess = _make_session("auto") |
| 217 | mock_common_state.store.list.return_value = {} # empty -> auto-create |
| 218 | mock_common_state.store.get.return_value = sess |
| 219 | new = mocker.patch("colab_cli.commands.session.new") |
| 220 | mocker.patch.object(ssh_module, "_resolve_pubkey", return_value="pk") |
| 221 | mocker.patch.object(ssh_module, "_run_interactive_ssh", return_value=0) |
| 222 | |
| 223 | result = runner.invoke(app, ["ssh", "--gpu", "T4", "--tpu", "v5e1"]) |
| 224 | assert result.exit_code == 0 |
| 225 | new.assert_called_once() |
| 226 | assert new.call_args.kwargs.get("gpu") == "T4" |
| 227 | assert new.call_args.kwargs.get("tpu") == "v5e1" |
| 228 | |
| 229 | |
| 230 | # --- F. --rm teardown is idempotent (signal path + finally path) ------------- |
| 231 | |
| 232 | |
| 233 | def test_proxy_mode_rm_teardown_idempotent(mock_common_state, mocker): |
| 234 | """If a signal fires mid-bridge AND the finally clean-close path runs, the |
| 235 | session is stopped exactly once (the `done` guard).""" |
| 236 | import signal as _signal |
| 237 | |
| 238 | sess = _make_session("colab-ephem") |
| 239 | mock_common_state.store.get.return_value = sess |
| 240 | mock_common_state.resolve_session.return_value = "colab-ephem" |
| 241 | |
| 242 | handlers = {} |
| 243 | mocker.patch( |
| 244 | "signal.signal", |
| 245 | side_effect=lambda sig, h: handlers.__setitem__(sig, h), |
| 246 | ) |
| 247 | mocker.patch("os._exit") # keep the handler from killing the test process |
| 248 | stop = mocker.patch("colab_cli.commands.session.stop") |
| 249 | mocker.patch.object(ssh_module, "_resolve_pubkey", return_value="pk") |
| 250 | mocker.patch.object( |
| 251 | ssh_module, "_connect_websocket", return_value=MagicMock() |
| 252 | ) |
| 253 | |
| 254 | def bridge_then_hup(ws): |
| 255 | handlers[_signal.SIGHUP](_signal.SIGHUP, None) # OpenSSH HUPs us |
| 256 | return 0 |
| 257 | |
| 258 | mocker.patch.object( |
| 259 | ssh_module, "_bridge_proxy_mode", side_effect=bridge_then_hup |
| 260 | ) |
| 261 | |
| 262 | result = runner.invoke( |
| 263 | app, ["ssh", "--proxy-mode", "-s", "colab-ephem", "--rm"] |
| 264 | ) |
| 265 | assert result.exit_code == 0 |
| 266 | stop.assert_called_once_with(session="colab-ephem") |
| 267 | |
| 268 | |
| 269 | # --- G. --proxy-mode --rm stops a *reused* session too ----------------------- |
| 270 | |
| 271 | |
| 272 | def test_proxy_mode_rm_stops_reused_session(mock_common_state, mocker): |
| 273 | """proxy-mode --rm stops the bridged session on disconnect even when it |
| 274 | already existed (ephemeral ~/.ssh/config host).""" |
| 275 | sess = _make_session("colab") |
| 276 | mock_common_state.store.get.return_value = sess # already exists |
| 277 | mock_common_state.resolve_session.return_value = "colab" |
| 278 | mocker.patch("signal.signal") |
| 279 | stop = mocker.patch("colab_cli.commands.session.stop") |
| 280 | _patch_proxy(mocker) |
| 281 | |
| 282 | result = runner.invoke(app, ["ssh", "--proxy-mode", "-s", "colab", "--rm"]) |
| 283 | assert result.exit_code == 0 |
| 284 | stop.assert_called_once_with(session="colab") |