@samitouri / QOSami-colab-cli
Code Commits Issues Pulls Sessions CI/CD
main
main
py 158 lines 5.91 KB
Raw
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, mock_open, patch
16
17 import pytest
18
19 from colab_cli.auth import (
20 REMOTE_REDIRECT_URI,
21 TOKEN_CONFIG_PATH,
22 AuthProvider,
23 get_credentials,
24 )
25
26
27 @pytest.fixture
28 def mock_deps(mocker):
29 m_exists = mocker.patch("os.path.exists")
30 m_makedirs = mocker.patch("os.makedirs")
31 m_creds_cls = mocker.patch("colab_cli.auth.Credentials")
32 m_flow_cls = mocker.patch("colab_cli.auth.InstalledAppFlow")
33 m_request = mocker.patch("colab_cli.auth.Request")
34 m_session = mocker.patch("colab_cli.auth.requests.AuthorizedSession")
35 m_resources = mocker.patch("colab_cli.auth.resources")
36
37 # By default, pretend oauth config doesn't exist anywhere
38 m_exists.return_value = False
39 m_resources.files.return_value.joinpath.return_value.is_file.return_value = False
40
41 return {
42 "exists": m_exists,
43 "makedirs": m_makedirs,
44 "creds_cls": m_creds_cls,
45 "flow_cls": m_flow_cls,
46 "request": m_request,
47 "session": m_session,
48 "resources": m_resources,
49 }
50
51
52 def test_get_credentials_no_config(mock_deps):
53 with pytest.raises(
54 FileNotFoundError,
55 match="Client OAuth config not found.*and no inlined config available",
56 ):
57 get_credentials("missing_config.json", provider=AuthProvider.OAUTH2)
58
59
60 def test_get_credentials_valid_token(mock_deps):
61 # Setup token exists
62 def exists_side_effect(path):
63 return path in ["dummy_config.json", TOKEN_CONFIG_PATH]
64
65 mock_deps["exists"].side_effect = exists_side_effect
66
67 # Valid creds
68 mock_creds = MagicMock()
69 mock_creds.valid = True
70 mock_deps["creds_cls"].from_authorized_user_file.return_value = mock_creds
71
72 # Mock open for config
73 m_open = mock_open(read_data='{"web":{"client_id":"id"}}')
74 with patch("builtins.open", m_open):
75 res = get_credentials("dummy_config.json", provider=AuthProvider.OAUTH2)
76
77 mock_deps["creds_cls"].from_authorized_user_file.assert_called_once()
78 mock_deps["session"].assert_called_once_with(mock_creds)
79 assert res == mock_deps["session"].return_value
80
81
82 def test_get_credentials_expired_token_refresh(mock_deps):
83 def exists_side_effect(path):
84 return path in ["dummy_config.json", TOKEN_CONFIG_PATH]
85
86 mock_deps["exists"].side_effect = exists_side_effect
87
88 mock_creds = MagicMock()
89 mock_creds.valid = False
90 mock_creds.expired = True
91 mock_creds.refresh_token = "some_token"
92 mock_creds.to_json.return_value = '{"token":"refreshed"}'
93 mock_deps["creds_cls"].from_authorized_user_file.return_value = mock_creds
94
95 m_open = mock_open(read_data='{"web":{"client_id":"id"}}')
96 with patch("builtins.open", m_open):
97 res = get_credentials("dummy_config.json", provider=AuthProvider.OAUTH2)
98
99 mock_creds.refresh.assert_called_once()
100 m_open.assert_any_call(TOKEN_CONFIG_PATH, "w")
101 assert res == mock_deps["session"].return_value
102
103
104 def test_get_credentials_no_token(mock_deps, mocker):
105 """With no cached token, the remote copy-paste flow runs and exchanges code."""
106 mock_deps["exists"].side_effect = lambda path: path == "dummy_config.json"
107
108 mock_flow = MagicMock()
109 mock_creds_new = MagicMock()
110 mock_creds_new.to_json.return_value = '{"token":"new"}'
111 mock_flow.authorization_url.return_value = ("https://auth.example/url", "state")
112 mock_flow.credentials = mock_creds_new
113 mock_deps["flow_cls"].from_client_config.return_value = mock_flow
114
115 # User pastes the authorization code at the prompt.
116 mocker.patch("colab_cli.auth.input", create=True, return_value="pasted-code")
117
118 m_open = mock_open(read_data='{"web":{"client_id":"id"}}')
119 with patch("builtins.open", m_open):
120 get_credentials("dummy_config.json", provider=AuthProvider.OAUTH2)
121
122 mock_deps["flow_cls"].from_client_config.assert_called_once()
123 # No localhost server should ever be started.
124 mock_flow.run_local_server.assert_not_called()
125 # Remote flow: OOB-free redirect + token_usage=remote consent param.
126 assert mock_flow.redirect_uri == REMOTE_REDIRECT_URI
127 _, kwargs = mock_flow.authorization_url.call_args
128 assert kwargs.get("token_usage") == "remote"
129 # The pasted code is exchanged for a token.
130 mock_flow.fetch_token.assert_called_once_with(code="pasted-code")
131
132
133 def test_remote_redirect_is_not_oob():
134 """Guard against regressing to the dead OOB redirect URI."""
135 assert REMOTE_REDIRECT_URI.startswith("https://")
136 assert "oob" not in REMOTE_REDIRECT_URI
137
138
139 def test_get_credentials_fallback_config(mock_deps):
140 # Setup: config_path doesn't exist, but fallback file does
141 mock_deps["exists"].return_value = False
142 m_file = mock_deps["resources"].files.return_value.joinpath.return_value
143 m_file.is_file.return_value = True
144 m_file.read_text.return_value = '{"installed":{"client_id":"fallback_id"}}'
145
146 # Valid creds in token
147 mock_deps["exists"].side_effect = lambda path: path == TOKEN_CONFIG_PATH
148 mock_creds = MagicMock()
149 mock_creds.valid = True
150 mock_deps["creds_cls"].from_authorized_user_file.return_value = mock_creds
151
152 res = get_credentials("missing_config.json", provider=AuthProvider.OAUTH2)
153
154 mock_deps["resources"].files.assert_called_once_with("colab_cli")
155 m_file.is_file.assert_called_once()
156 m_file.read_text.assert_called_once()
157 mock_deps["creds_cls"].from_authorized_user_file.assert_called_once()
158 assert res == mock_deps["session"].return_value