| 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 | """`colab ssh` interactive shell: lands in /content and forwards --identity. |
| 16 | |
| 17 | An SSH login lands in root's home (/root); Colab users expect /content (where |
| 18 | notebooks + uploads live). `_run_interactive_ssh` forces a PTY and runs a remote |
| 19 | command that `cd`s to /content before exec'ing the login shell, and threads |
| 20 | `--identity` through to both the outer `ssh -i` and the inner proxy command. |
| 21 | """ |
| 22 | |
| 23 | from unittest.mock import MagicMock |
| 24 | |
| 25 | from colab_cli.commands import ssh as ssh_module |
| 26 | import pytest |
| 27 | |
| 28 | |
| 29 | def _make_session( |
| 30 | name: str = "s1", |
| 31 | url: str = "https://abc.colab.googleusercontent.com", |
| 32 | token: str = "TOK", |
| 33 | endpoint: str = "ep", |
| 34 | ): |
| 35 | s = MagicMock() |
| 36 | s.name = name |
| 37 | s.url = url |
| 38 | s.token = token |
| 39 | s.endpoint = endpoint |
| 40 | return s |
| 41 | |
| 42 | |
| 43 | def test_default_remote_dir_is_content(): |
| 44 | assert ssh_module._DEFAULT_REMOTE_DIR == "/content" |
| 45 | |
| 46 | |
| 47 | @pytest.mark.parametrize( |
| 48 | "identity", |
| 49 | [None, "/home/u/.ssh/id_ed25519"], |
| 50 | ids=["no-identity", "with-identity"], |
| 51 | ) |
| 52 | def test_interactive_ssh_builds_args(mocker, identity): |
| 53 | call = mocker.patch("subprocess.call", return_value=0) |
| 54 | rc = ssh_module._run_interactive_ssh(_make_session(), identity) |
| 55 | assert rc == 0 |
| 56 | |
| 57 | args = call.call_args.args[0] |
| 58 | assert "-t" in args # PTY forced so the exec'd shell is interactive |
| 59 | assert ssh_module._SSH_HOST in args |
| 60 | # the host must precede the remote command (which is the final element) |
| 61 | assert args.index(ssh_module._SSH_HOST) < len(args) - 1 |
| 62 | |
| 63 | remote = args[-1] |
| 64 | assert "cd /content" in remote |
| 65 | assert "exec" in remote # execs a shell after the cd |
| 66 | # a missing /content must not abort the shell (stderr suppressed). |
| 67 | assert "2>/dev/null" in remote |
| 68 | |
| 69 | joined = " ".join(args) |
| 70 | if identity: |
| 71 | assert "-i" in args # outer ssh identity |
| 72 | assert "--identity" in joined # inner proxy-mode identity |
| 73 | else: |
| 74 | assert "-i" not in args |
| 75 | assert "--identity" not in joined |