@samitouri / QOSami-colab-cli-2
Code Commits Issues Pulls Sessions CI/CD
main
main
py 172 lines 6.66 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 """Tests for the developer-only `colab whoami` command.
16
17 This is a debugging aid that resolves the active credentials, mints an access
18 token, and queries Google's tokeninfo endpoint to print the human-readable
19 identity (email), grant scopes, and expiry of whatever the CLI is about to
20 send to colab.research.google.com / colab.pa.googleapis.com.
21 """
22
23 import sys
24 from unittest.mock import MagicMock, patch
25
26 import pytest
27
28 from colab_cli.cli import main
29 from colab_cli.auth import AuthProvider
30
31
32 # Unambiguously-fake placeholders so credential-scanner pre-commit hooks
33 # don't false-positive on `ya29.*` strings. The whoami code only ever uses
34 # these as opaque payloads passed to a mocked urllib.request.urlopen.
35 _FAKE_TOKEN = "TEST-TOKEN-PLACEHOLDER"
36
37
38 def _fake_creds(token: str = _FAKE_TOKEN):
39 """Build a credentials-like mock: has .token and .refresh()."""
40 creds = MagicMock()
41 creds.token = token
42 creds.refresh = MagicMock()
43 return creds
44
45
46 def _fake_authed_session(token: str = _FAKE_TOKEN):
47 """Mimic google.auth.transport.requests.AuthorizedSession enough for whoami."""
48 sess = MagicMock()
49 sess.credentials = _fake_creds(token)
50 return sess
51
52
53 def test_whoami_prints_human_readable_summary(mock_common_state, capsys):
54 """Default invocation should fetch the token, hit tokeninfo, and print
55 a labelled summary including email, the active auth provider, scopes
56 (one per line), and an Expires line."""
57 mock_common_state.auth_provider = AuthProvider.ADC
58
59 fake_response = MagicMock()
60 fake_response.status_code = 200
61 fake_response.json.return_value = {
62 "email": "user@example.com",
63 "scope": (
64 "https://www.googleapis.com/auth/userinfo.email "
65 "https://www.googleapis.com/auth/colaboratory"
66 ),
67 "expires_in": "2847",
68 "audience": "32555940559.apps.googleusercontent.com",
69 }
70
71 with patch("colab_cli.auth.get_credentials", return_value=_fake_authed_session()):
72 with patch("urllib.request.urlopen") as mock_urlopen:
73 cm = MagicMock()
74 cm.__enter__ = MagicMock(return_value=cm)
75 cm.__exit__ = MagicMock(return_value=False)
76 cm.read.return_value = b'{"email":"user@example.com","scope":"https://www.googleapis.com/auth/userinfo.email https://www.googleapis.com/auth/colaboratory","expires_in":"2847","audience":"32555940559.apps.googleusercontent.com"}'
77 cm.status = 200
78 mock_urlopen.return_value = cm
79
80 with patch.object(sys, "argv", ["colab", "--auth=adc", "whoami"]):
81 with pytest.raises(SystemExit) as error:
82 main()
83 assert error.value.code == 0
84
85 out = capsys.readouterr().out
86 assert "user@example.com" in out
87 assert "adc" in out.lower()
88 assert "userinfo.email" in out
89 assert "colaboratory" in out
90 # Expiry should be rendered in a human form (minutes), not raw seconds.
91 assert "47m" in out or "47 min" in out
92
93
94 def test_whoami_handles_tokeninfo_error_gracefully(mock_common_state, capsys):
95 """If tokeninfo returns 4xx (e.g. expired/revoked token), whoami should
96 print an error and exit non-zero rather than blowing up with a stack trace.
97 Developers reading the message should be able to tell what happened.
98 """
99 mock_common_state.auth_provider = AuthProvider.ADC
100
101 with patch("colab_cli.auth.get_credentials", return_value=_fake_authed_session()):
102 with patch("urllib.request.urlopen") as mock_urlopen:
103 import urllib.error
104
105 mock_urlopen.side_effect = urllib.error.HTTPError(
106 url="https://oauth2.googleapis.com/tokeninfo",
107 code=400,
108 msg="Bad Request",
109 hdrs=None,
110 fp=None,
111 )
112
113 with patch.object(sys, "argv", ["colab", "--auth=adc", "whoami"]):
114 with pytest.raises(SystemExit) as error:
115 main()
116 assert error.value.code != 0
117
118 captured = capsys.readouterr()
119 assert "tokeninfo" in (captured.out + captured.err).lower() or "400" in (
120 captured.out + captured.err
121 )
122
123
124 def test_whoami_is_hidden_from_top_level_help(mock_common_state, capsys):
125 """`colab --help` should not list `whoami` (it's a developer tool).
126 Also asserts that `colab whoami --help` still works (the command is
127 hidden, not removed)."""
128 # `--help` exits 0
129 with patch.object(sys, "argv", ["colab", "--help"]):
130 with pytest.raises(SystemExit) as error:
131 main()
132 assert error.value.code == 0
133 out = capsys.readouterr().out
134 assert "whoami" not in out, (
135 f"`whoami` should be hidden from `colab --help`, but appeared in:\n{out}"
136 )
137
138 # Confirm `colab whoami --help` is still reachable.
139 with patch.object(sys, "argv", ["colab", "whoami", "--help"]):
140 with pytest.raises(SystemExit) as error:
141 main()
142 assert error.value.code == 0
143 out2 = capsys.readouterr().out
144 assert "whoami" in out2.lower(), (
145 f"`colab whoami --help` should describe the command, got:\n{out2}"
146 )
147
148
149 def test_whoami_refreshes_credentials_before_reading_token(mock_common_state):
150 """Some credentials (ADC service account, GCE) lazy-mint the token only
151 when refresh() is called. whoami must call refresh() before reading
152 creds.token, otherwise creds.token may be None even for valid creds.
153 """
154 mock_common_state.auth_provider = AuthProvider.ADC
155 sess = _fake_authed_session(token="TEST-TOKEN-AFTER-REFRESH")
156
157 with patch("colab_cli.auth.get_credentials", return_value=sess):
158 with patch("urllib.request.urlopen") as mock_urlopen:
159 cm = MagicMock()
160 cm.__enter__ = MagicMock(return_value=cm)
161 cm.__exit__ = MagicMock(return_value=False)
162 cm.read.return_value = (
163 b'{"email":"x@y.com","scope":"a b","expires_in":"60"}'
164 )
165 cm.status = 200
166 mock_urlopen.return_value = cm
167
168 with patch.object(sys, "argv", ["colab", "--auth=adc", "whoami"]):
169 with pytest.raises(SystemExit):
170 main()
171
172 sess.credentials.refresh.assert_called_once()