| 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 patch |
| 16 | import pytest |
| 17 | from typer.testing import CliRunner |
| 18 | from colab_cli.cli import app |
| 19 | from colab_cli.state import SessionState |
| 20 | |
| 21 | runner = CliRunner() |
| 22 | |
| 23 | |
| 24 | @pytest.fixture |
| 25 | def mock_session(): |
| 26 | return SessionState( |
| 27 | name="test-session", |
| 28 | token="test-token", |
| 29 | url="https://test.url", |
| 30 | endpoint="e1", |
| 31 | ) |
| 32 | |
| 33 | |
| 34 | @patch("colab_cli.commands.automation.ColabRuntime") |
| 35 | @patch("colab_cli.common.state") |
| 36 | def test_cli_auth(mock_state, mock_runtime_class, mock_session): |
| 37 | mock_state.store.get.return_value = mock_session |
| 38 | mock_state.resolve_session.return_value = "test-session" |
| 39 | |
| 40 | mock_runtime = mock_runtime_class.return_value |
| 41 | mock_runtime.execute_code.return_value = [{"text": "Success"}] |
| 42 | |
| 43 | result = runner.invoke(app, ["auth", "-s", "test-session"]) |
| 44 | assert result.exit_code == 0 |
| 45 | |
| 46 | assert mock_session.last_execution[0] == "automation:auth" |
| 47 | assert mock_session.last_execution[1] is None |
| 48 | assert mock_session.last_execution[2] is not None |
| 49 | mock_state.store.add.assert_called_with(mock_session) |
| 50 | |
| 51 | # Verify ColabRuntime was invoked with the correct code |
| 52 | mock_runtime.execute_code.assert_called_once() |
| 53 | called_code = mock_runtime.execute_code.call_args[0][0] |
| 54 | |
| 55 | assert "os.environ['USE_AUTH_EPHEM'] = '0'" in called_code |
| 56 | assert "auth.authenticate_user()" in called_code |
| 57 | |
| 58 | |
| 59 | @patch("colab_cli.commands.automation.ColabRuntime") |
| 60 | @patch("colab_cli.common.state") |
| 61 | def test_cli_install(mock_state, mock_runtime_class, mock_session): |
| 62 | mock_state.store.get.return_value = mock_session |
| 63 | mock_state.resolve_session.return_value = "test-session" |
| 64 | |
| 65 | mock_runtime = mock_runtime_class.return_value |
| 66 | mock_runtime.execute_code.return_value = [{"text": "Installed"}] |
| 67 | |
| 68 | result = runner.invoke(app, ["install", "-s", "test-session", "pandas", "numpy"]) |
| 69 | assert result.exit_code == 0 |
| 70 | assert mock_session.last_execution[0] == "automation:install" |
| 71 | assert mock_session.last_execution[2] is not None |
| 72 | mock_state.store.add.assert_called_with(mock_session) |
| 73 | |
| 74 | mock_runtime.execute_code.assert_called_once() |
| 75 | called_code = mock_runtime.execute_code.call_args[0][0] |
| 76 | |
| 77 | assert "subprocess" in called_code |
| 78 | assert "pip" in called_code |
| 79 | assert "pandas" in called_code |
| 80 | assert "numpy" in called_code |
| 81 | |
| 82 | |
| 83 | @patch("colab_cli.commands.automation.ColabRuntime") |
| 84 | @patch("colab_cli.common.state") |
| 85 | def test_cli_drivemount(mock_state, mock_runtime_class, mock_session): |
| 86 | mock_state.store.get.return_value = mock_session |
| 87 | mock_state.resolve_session.return_value = "test-session" |
| 88 | |
| 89 | mock_runtime = mock_runtime_class.return_value |
| 90 | mock_runtime.execute_code.return_value = [{"text": "Mounted"}] |
| 91 | |
| 92 | result = runner.invoke(app, ["drivemount", "-s", "test-session", "/foo/bar"]) |
| 93 | assert result.exit_code == 0 |
| 94 | |
| 95 | # Verify ColabRuntime was invoked with the correct code |
| 96 | mock_runtime.execute_code.assert_called_once() |
| 97 | called_code = mock_runtime.execute_code.call_args[0][0] |
| 98 | |
| 99 | assert "drive.mount('/foo/bar')" in called_code |
| 100 | assert mock_runtime.colab_request_hook is not None |
| 101 | # Drivemount waits for the user to OAuth in their browser; the kernel |
| 102 | # goes silent during that wait and the default 10s execute() timeout |
| 103 | # would raise TimeoutError mid-flow. Insist on a generous timeout |
| 104 | # (>= 5 minutes) being forwarded to runtime.execute_code. |
| 105 | _, kwargs = mock_runtime.execute_code.call_args |
| 106 | assert kwargs.get("timeout") is not None and kwargs["timeout"] >= 300 |
| 107 | |
| 108 | |
| 109 | @patch("colab_cli.commands.automation.ColabRuntime") |
| 110 | @patch("colab_cli.common.state") |
| 111 | def test_cli_auth_uses_long_timeout(mock_state, mock_runtime_class, mock_session): |
| 112 | """`colab auth` walks the user through a paste-the-code flow that |
| 113 | routinely takes >10s, so it must pass a generous timeout to |
| 114 | runtime.execute_code or the call will TimeoutError mid-flow.""" |
| 115 | mock_state.store.get.return_value = mock_session |
| 116 | mock_state.resolve_session.return_value = "test-session" |
| 117 | |
| 118 | mock_runtime = mock_runtime_class.return_value |
| 119 | mock_runtime.execute_code.return_value = [{"text": "Authenticated"}] |
| 120 | |
| 121 | result = runner.invoke(app, ["auth", "-s", "test-session"]) |
| 122 | assert result.exit_code == 0 |
| 123 | |
| 124 | _, kwargs = mock_runtime.execute_code.call_args |
| 125 | assert kwargs.get("timeout") is not None and kwargs["timeout"] >= 300 |