| 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 `colab url`: print a browser URL that connects the Colab |
| 16 | frontend to an existing colab-cli session. |
| 17 | |
| 18 | URL format: |
| 19 | |
| 20 | https://<host>/notebooks/empty.ipynb?dbu=<urlencoded-/tun/m/endpoint>#datalabBackendUrl=<host>/tun/m/<endpoint> |
| 21 | |
| 22 | Two backend-URL signals are embedded: |
| 23 | |
| 24 | - `?dbu=<urlencoded path>` -- the Colab frontend's |
| 25 | `datalab_backend_url` development query flag. The frontend resolves |
| 26 | the value against `window.location.origin` and attaches the kernel |
| 27 | to the supplied `/tun/m/<endpoint>` path instead of allocating a |
| 28 | fresh VM. |
| 29 | |
| 30 | - `#datalabBackendUrl=<full URL>` -- the hash-fragment form. Some |
| 31 | Colab frontend code paths consult this first and ignore `dbu`, so we |
| 32 | emit both for robustness. The fragment value is a FULL URL (with |
| 33 | scheme + host) and is intentionally NOT URL-encoded -- browsers do |
| 34 | not decode fragment values before passing them to page JS, and |
| 35 | Colab's hash parser expects the raw string. |
| 36 | |
| 37 | The fragment's host always matches the page origin (`--host`), so |
| 38 | same-origin enforcement in the frontend doesn't block the connection |
| 39 | and sandbox/dev users get a sandbox fragment automatically. |
| 40 | """ |
| 41 | |
| 42 | from unittest.mock import MagicMock, patch |
| 43 | from urllib.parse import parse_qs, quote, urlparse |
| 44 | |
| 45 | from typer.testing import CliRunner |
| 46 | |
| 47 | from colab_cli.cli import app |
| 48 | |
| 49 | runner = CliRunner() |
| 50 | |
| 51 | |
| 52 | def _make_session(name: str = "s1", endpoint: str = "abc123def"): |
| 53 | s = MagicMock() |
| 54 | s.name = name |
| 55 | s.endpoint = endpoint |
| 56 | return s |
| 57 | |
| 58 | |
| 59 | def _parse_url_output(output: str) -> str: |
| 60 | """Pull the single URL line out of `colab url` output.""" |
| 61 | candidates = [line.strip() for line in output.splitlines() if "dbu=" in line] |
| 62 | assert len(candidates) == 1, ( |
| 63 | f"Expected exactly one URL line containing 'dbu=', got {candidates!r}" |
| 64 | ) |
| 65 | return candidates[0] |
| 66 | |
| 67 | |
| 68 | def test_url_explicit_session(mock_common_state): |
| 69 | """`colab url -s NAME` prints the connect URL for that session. |
| 70 | |
| 71 | Format: ``https://<host>/notebooks/empty.ipynb?dbu=<urlencoded path>#datalabBackendUrl=<host>/tun/m/<endpoint>``. |
| 72 | The path must land on `empty.ipynb` so the user sees a usable notebook |
| 73 | UI; the `dbu` query param tells the frontend to skip /tun/m/assign and |
| 74 | attach to our existing endpoint; the `#datalabBackendUrl=` fragment |
| 75 | is the alternative signal some frontend code paths consult. |
| 76 | """ |
| 77 | s = _make_session(name="my-sess", endpoint="ep-XYZ") |
| 78 | mock_common_state.store.get.return_value = s |
| 79 | mock_common_state.resolve_session.return_value = "my-sess" |
| 80 | |
| 81 | result = runner.invoke(app, ["url", "-s", "my-sess"]) |
| 82 | |
| 83 | assert result.exit_code == 0, result.output |
| 84 | url = _parse_url_output(result.output) |
| 85 | |
| 86 | parsed = urlparse(url) |
| 87 | assert parsed.scheme == "https" |
| 88 | assert parsed.netloc == "colab.research.google.com" |
| 89 | assert parsed.path == "/notebooks/empty.ipynb" |
| 90 | |
| 91 | # `dbu` must be the URL-encoded path `/tun/m/<endpoint>`. We assert on |
| 92 | # the decoded form rather than the raw encoding to keep the test robust |
| 93 | # to which characters the encoder happens to escape (e.g. `/` may or |
| 94 | # may not be escaped depending on `safe=`); what matters is round-trip |
| 95 | # decoding produces the right backend path. |
| 96 | qs = parse_qs(parsed.query) |
| 97 | assert qs.get("dbu") == ["/tun/m/ep-XYZ"] |
| 98 | |
| 99 | # And we DO actually URL-encode the slashes so the value survives any |
| 100 | # downstream re-parsing that treats the query string non-strictly. |
| 101 | assert "dbu=%2Ftun%2Fm%2Fep-XYZ" in url |
| 102 | |
| 103 | # Fragment: raw (not URL-encoded), full URL form, host matches page origin. |
| 104 | assert ( |
| 105 | parsed.fragment |
| 106 | == "datalabBackendUrl=https://colab.research.google.com/tun/m/ep-XYZ" |
| 107 | ) |
| 108 | |
| 109 | |
| 110 | def test_url_resolves_unique_session(mock_common_state): |
| 111 | """`colab url` (no -s) uses the unique-session resolution path.""" |
| 112 | s = _make_session(name="only-sess", endpoint="solo-EP") |
| 113 | mock_common_state.store.get.return_value = s |
| 114 | mock_common_state.resolve_session.return_value = "only-sess" |
| 115 | |
| 116 | result = runner.invoke(app, ["url"]) |
| 117 | |
| 118 | assert result.exit_code == 0, result.output |
| 119 | url = _parse_url_output(result.output) |
| 120 | assert "%2Ftun%2Fm%2Fsolo-EP" in url |
| 121 | # Resolution went through the shared helper, not by hardcoding the name. |
| 122 | mock_common_state.resolve_session.assert_called_once_with(None) |
| 123 | |
| 124 | |
| 125 | def test_url_session_not_found(mock_common_state): |
| 126 | """If the resolved session has no local state, exit non-zero with a clear |
| 127 | message rather than printing a malformed URL.""" |
| 128 | mock_common_state.resolve_session.return_value = "ghost" |
| 129 | mock_common_state.store.get.return_value = None |
| 130 | |
| 131 | result = runner.invoke(app, ["url", "-s", "ghost"]) |
| 132 | |
| 133 | assert result.exit_code != 0 |
| 134 | assert "ghost" in result.output |
| 135 | assert "not found" in result.output.lower() |
| 136 | |
| 137 | |
| 138 | def test_url_custom_host(mock_common_state): |
| 139 | """`--host` overrides the default frontend host AND the host used in |
| 140 | the `#datalabBackendUrl=` fragment. |
| 141 | |
| 142 | `dbu` itself is a path-only value (resolved against |
| 143 | `window.location.origin` in the frontend), so the host swap only |
| 144 | affects the page origin, not the embedded backend path. But the |
| 145 | fragment carries a full URL, and the Colab frontend enforces |
| 146 | same-origin between page and embedded backend URL -- so the fragment |
| 147 | host MUST match `--host` for the swap to work end-to-end. |
| 148 | """ |
| 149 | s = _make_session(endpoint="ep1") |
| 150 | mock_common_state.store.get.return_value = s |
| 151 | mock_common_state.resolve_session.return_value = "s1" |
| 152 | |
| 153 | result = runner.invoke( |
| 154 | app, ["url", "-s", "s1", "--host", "https://colab.sandbox.google.com"] |
| 155 | ) |
| 156 | |
| 157 | assert result.exit_code == 0, result.output |
| 158 | url = _parse_url_output(result.output) |
| 159 | parsed = urlparse(url) |
| 160 | assert parsed.netloc == "colab.sandbox.google.com" |
| 161 | assert parsed.path == "/notebooks/empty.ipynb" |
| 162 | assert parse_qs(parsed.query).get("dbu") == ["/tun/m/ep1"] |
| 163 | # Fragment host tracks --host (NOT pinned to research.google.com). |
| 164 | assert ( |
| 165 | parsed.fragment |
| 166 | == "datalabBackendUrl=https://colab.sandbox.google.com/tun/m/ep1" |
| 167 | ) |
| 168 | |
| 169 | |
| 170 | def test_url_host_normalises_trailing_slash(mock_common_state): |
| 171 | """`--host https://example.com/` (with trailing slash) must not produce |
| 172 | a double slash anywhere -- not before `/notebooks/empty.ipynb` in the |
| 173 | page URL, AND not before `/tun/m/...` in the fragment value.""" |
| 174 | s = _make_session(endpoint="ep2") |
| 175 | mock_common_state.store.get.return_value = s |
| 176 | mock_common_state.resolve_session.return_value = "s1" |
| 177 | |
| 178 | result = runner.invoke( |
| 179 | app, ["url", "-s", "s1", "--host", "https://colab.research.google.com/"] |
| 180 | ) |
| 181 | |
| 182 | assert result.exit_code == 0 |
| 183 | assert "https://colab.research.google.com//notebooks/" not in result.output |
| 184 | assert "https://colab.research.google.com/notebooks/empty.ipynb" in result.output |
| 185 | # Same guarantee for the fragment URL. |
| 186 | assert "https://colab.research.google.com//tun/" not in result.output |
| 187 | assert ( |
| 188 | "datalabBackendUrl=https://colab.research.google.com/tun/m/ep2" in result.output |
| 189 | ) |
| 190 | |
| 191 | |
| 192 | def test_url_endpoint_with_special_chars_is_encoded(mock_common_state): |
| 193 | """Endpoints are opaque server-issued IDs but we should not assume their |
| 194 | character set. Anything outside the unreserved URL set must be escaped |
| 195 | so the frontend's `new URL(...)` parser sees the intended path.""" |
| 196 | # Intentionally include characters that MUST be escaped if they ever |
| 197 | # appeared in an endpoint (e.g. `&`, `?`, `#`, space, `=`). |
| 198 | s = _make_session(endpoint="weird ep&?=#") |
| 199 | mock_common_state.store.get.return_value = s |
| 200 | mock_common_state.resolve_session.return_value = "s1" |
| 201 | |
| 202 | result = runner.invoke(app, ["url", "-s", "s1"]) |
| 203 | assert result.exit_code == 0, result.output |
| 204 | url = _parse_url_output(result.output) |
| 205 | parsed = urlparse(url) |
| 206 | |
| 207 | # The encoded endpoint must round-trip via the standard query parser. |
| 208 | assert parse_qs(parsed.query).get("dbu") == ["/tun/m/weird ep&?=#"] |
| 209 | # And the raw URL must contain the percent-encoded form (not the literal). |
| 210 | assert quote("weird ep&?=#", safe="") in url |
| 211 | |
| 212 | |
| 213 | def test_url_open_flag_launches_browser(mock_common_state): |
| 214 | """`--open` calls webbrowser.open() with the same URL it printed.""" |
| 215 | s = _make_session(endpoint="ep-OPEN") |
| 216 | mock_common_state.store.get.return_value = s |
| 217 | mock_common_state.resolve_session.return_value = "s1" |
| 218 | |
| 219 | with patch("webbrowser.open") as mock_open: |
| 220 | result = runner.invoke(app, ["url", "-s", "s1", "--open"]) |
| 221 | |
| 222 | assert result.exit_code == 0, result.output |
| 223 | mock_open.assert_called_once() |
| 224 | opened_url = mock_open.call_args[0][0] |
| 225 | assert "dbu=" in opened_url |
| 226 | assert "%2Ftun%2Fm%2Fep-OPEN" in opened_url |
| 227 | # And the URL was also printed (so users see what was opened, and |
| 228 | # piping still works). |
| 229 | assert opened_url in result.output |
| 230 | |
| 231 | |
| 232 | def test_url_no_open_by_default(mock_common_state): |
| 233 | """Default behaviour: print only, do NOT auto-open the browser. This keeps |
| 234 | the command pipeable (`colab url | xclip`, `colab url | pbcopy`, etc.).""" |
| 235 | s = _make_session(endpoint="ep3") |
| 236 | mock_common_state.store.get.return_value = s |
| 237 | mock_common_state.resolve_session.return_value = "s1" |
| 238 | |
| 239 | with patch("webbrowser.open") as mock_open: |
| 240 | result = runner.invoke(app, ["url", "-s", "s1"]) |
| 241 | |
| 242 | assert result.exit_code == 0 |
| 243 | mock_open.assert_not_called() |
| 244 | |
| 245 | |
| 246 | def test_url_fragment_is_not_url_encoded(mock_common_state): |
| 247 | """The `#datalabBackendUrl=...` fragment value is a full URL and must |
| 248 | be embedded raw (no percent-encoding). Browsers do not decode the |
| 249 | fragment before passing `location.hash` to page JS, and the Colab |
| 250 | parser expects to call `new URL(rawString)` directly. If we encoded |
| 251 | `:` -> `%3A` or `/` -> `%2F` here, the parser would see |
| 252 | `https%3A%2F%2Fcolab...` and fail. |
| 253 | |
| 254 | Concretely we should see the literal `://` and unescaped `/` in the |
| 255 | fragment, NOT their percent-encoded counterparts. |
| 256 | """ |
| 257 | s = _make_session(endpoint="ep-RAW") |
| 258 | mock_common_state.store.get.return_value = s |
| 259 | mock_common_state.resolve_session.return_value = "s1" |
| 260 | |
| 261 | result = runner.invoke(app, ["url", "-s", "s1"]) |
| 262 | assert result.exit_code == 0, result.output |
| 263 | url = _parse_url_output(result.output) |
| 264 | parsed = urlparse(url) |
| 265 | |
| 266 | # The fragment must contain the literal scheme + slashes... |
| 267 | assert "datalabBackendUrl=https://" in url |
| 268 | assert "/tun/m/ep-RAW" in parsed.fragment |
| 269 | # ...and MUST NOT contain percent-encoded versions of `:`, `/`. |
| 270 | assert "%3A" not in parsed.fragment |
| 271 | assert "%2F" not in parsed.fragment |
| 272 | |
| 273 | |
| 274 | def test_url_both_signals_present(mock_common_state): |
| 275 | """Invariant: every printed URL has BOTH `?dbu=` and `#datalabBackendUrl=`. |
| 276 | |
| 277 | Either alone is unreliable across Colab frontend revisions; we emit |
| 278 | both so the frontend can use whichever it consults first. |
| 279 | """ |
| 280 | s = _make_session(endpoint="ep-BOTH") |
| 281 | mock_common_state.store.get.return_value = s |
| 282 | mock_common_state.resolve_session.return_value = "s1" |
| 283 | |
| 284 | result = runner.invoke(app, ["url", "-s", "s1"]) |
| 285 | assert result.exit_code == 0, result.output |
| 286 | url = _parse_url_output(result.output) |
| 287 | assert "?dbu=" in url |
| 288 | assert "#datalabBackendUrl=" in url |
| 289 | |
| 290 | |
| 291 | def test_url_open_flag_includes_fragment(mock_common_state): |
| 292 | """`--open` opens the SAME URL it printed, including the fragment. |
| 293 | |
| 294 | `webbrowser.open()` must receive the URL with the `#datalabBackendUrl=` |
| 295 | fragment intact -- otherwise the browser may attach to a fresh VM via |
| 296 | `/tun/m/assign` instead of our existing session. |
| 297 | """ |
| 298 | s = _make_session(endpoint="ep-OPEN2") |
| 299 | mock_common_state.store.get.return_value = s |
| 300 | mock_common_state.resolve_session.return_value = "s1" |
| 301 | |
| 302 | with patch("webbrowser.open") as mock_open: |
| 303 | result = runner.invoke(app, ["url", "-s", "s1", "--open"]) |
| 304 | |
| 305 | assert result.exit_code == 0, result.output |
| 306 | mock_open.assert_called_once() |
| 307 | opened_url = mock_open.call_args[0][0] |
| 308 | assert ( |
| 309 | "#datalabBackendUrl=https://colab.research.google.com/tun/m/ep-OPEN2" |
| 310 | in opened_url |
| 311 | ) |
| 312 | # And it's the same URL that got printed. |
| 313 | assert opened_url in result.output |
| 314 | |
| 315 | |
| 316 | def test_url_output_is_pipeable(mock_common_state): |
| 317 | """The printed URL line must be machine-parseable: a single line with no |
| 318 | leading `[colab]` chatter, so `colab url -s s1 | xclip` works. |
| 319 | """ |
| 320 | s = _make_session(endpoint="ep-PIPE") |
| 321 | mock_common_state.store.get.return_value = s |
| 322 | mock_common_state.resolve_session.return_value = "s1" |
| 323 | |
| 324 | result = runner.invoke(app, ["url", "-s", "s1"]) |
| 325 | |
| 326 | assert result.exit_code == 0 |
| 327 | url_lines = [line for line in result.output.splitlines() if "dbu=" in line] |
| 328 | assert len(url_lines) == 1, f"Expected exactly one URL line, got: {url_lines}" |
| 329 | assert not url_lines[0].lstrip().startswith("[colab]"), ( |
| 330 | f"URL line should not be prefixed with '[colab]' so it's pipeable: " |
| 331 | f"{url_lines[0]!r}" |
| 332 | ) |