diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index 5ed537c..ab6d188 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -421,6 +421,11 @@ def has_valid_databricks_auth(workspace: str, profile: str | None = None) -> boo def get_databricks_profiles() -> list[tuple[str, str]]: """Return [(host_url, profile_name), ...] from Databricks CLI profiles. + Each non-PAT profile is returned individually — duplicate hosts (multiple + profiles pointing at the same workspace) appear as separate entries so the + workspace picker can offer each profile by name. Order matches the CLI's + own ordering. + Returns ``[]`` on any failure (CLI missing, timeout, non-zero exit, JSON decode error). When ``UCODE_DEBUG=1`` each dropout path logs *why* the result was empty so a silently-disappearing workspace picker is @@ -446,8 +451,7 @@ def get_databricks_profiles() -> list[tuple[str, str]]: _debug("get_databricks_profiles", f"json decode error: {exc.msg}") return [] - # dict dedupes by host (first non-PAT profile wins). - out: dict[str, str] = {} + out: list[tuple[str, str]] = [] pat = 0 for p in profiles: host = (p.get("host") or "").rstrip("/") @@ -457,13 +461,13 @@ def get_databricks_profiles() -> list[tuple[str, str]]: if p.get("auth_type") == "pat": pat += 1 continue - out.setdefault(host, name) + out.append((host, name)) _debug( "get_databricks_profiles", f"returned={len(out)} total={len(profiles)} pat={pat}", ) - return list(out.items()) + return out def find_profile_name_for_host(workspace: str) -> str | None: diff --git a/src/ucode/ui.py b/src/ucode/ui.py index 4f40d1b..1d85152 100644 --- a/src/ucode/ui.py +++ b/src/ucode/ui.py @@ -181,23 +181,48 @@ def prompt_for_workspace( """Ask the user for a workspace URL, offering profiles as quick-select. `profiles` is a list of (host_url, profile_name) tuples. Caller fetches - them — `ui.py` stays Databricks-agnostic. Returns ``(url, profile_name)``; - profile_name is ``None`` when the user typed a URL manually. + them — `ui.py` stays Databricks-agnostic. Duplicate hosts (multiple + profiles pointing at the same workspace) are shown separately; the picker + returns the exact (host, profile_name) the user selected. Returns + ``(url, profile_name)``; profile_name is ``None`` when the user typed a + URL manually. """ console.print() console.print(Panel(description, title="ucode setup", style="bold blue", expand=False)) if profiles: - choices = [ - questionary.Choice(title=host, value=(host, profile_name)) - for host, profile_name in profiles + name_header = "Profile Name" + url_header = "Workspace URL" + # Clamp so a single very long profile name can't push the URL column + # off-screen on an 80-col terminal — questionary doesn't wrap row + # titles cleanly, and a wrapped row breaks the picker visually. + max_name_width = 40 + name_width = min( + max_name_width, + max(len(name_header), *(len(name) for _, name in profiles)), + ) + # Match the 2-char cursor gutter so the header line aligns with rows. + header_title = f" {name_header.ljust(name_width)} {url_header}" + choices: list[questionary.Choice | questionary.Separator] = [ + questionary.Separator(header_title) ] + for host, profile_name in profiles: + display_name = ( + profile_name + if len(profile_name) <= name_width + else profile_name[: name_width - 1] + "…" + ) + row_title = f"{display_name.ljust(name_width)} {host}" + # Value carries the full untruncated profile name so downstream + # `--profile` calls always use the real name, not the display form. + choices.append(questionary.Choice(title=row_title, value=(host, profile_name))) choices.append(questionary.Choice(title="Enter a different URL", value=None)) style = questionary.Style( [ ("highlighted", "fg:cyan bold"), ("pointer", "fg:cyan bold"), ("answer", "fg:cyan"), + ("separator", "fg:white bold"), ] ) choice = questionary.select( diff --git a/tests/test_cli.py b/tests/test_cli.py index e2c95e2..957b341 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -602,6 +602,39 @@ def test_unavailable_selected_tool_errors_before_configure(self, monkeypatch): with pytest.raises(RuntimeError, match="Codex"): cli_mod.configure_workspace_command(selected_tools=["claude", "codex"]) + def test_picker_selected_profile_flows_to_configure_shared_state(self, monkeypatch): + """Picker's (host, profile) tuple must reach configure_shared_state's + `profile` kwarg, otherwise downstream --profile calls fall back to + host-based resolution and silently pick the wrong profile.""" + import ucode.cli as cli_mod + + monkeypatch.setattr( + cli_mod, + "_prompt_for_configuration", + lambda tool=None: ("https://shared.cloud.databricks.com", "picked-profile"), + ) + captured: dict = {} + + def fake_configure_shared_state(workspace, profile=None, tools=None, force_login=False): + captured["workspace"] = workspace + captured["profile"] = profile + return {**MINIMAL_STATE, "workspace": workspace, "profile": profile} + + monkeypatch.setattr(cli_mod, "configure_shared_state", fake_configure_shared_state) + monkeypatch.setattr(cli_mod, "save_state", lambda state: None) + monkeypatch.setattr(cli_mod, "check_gateway_endpoint", lambda state, tool: True) + monkeypatch.setattr(cli_mod, "prompt_for_tools", lambda available: ["claude"]) + monkeypatch.setattr(cli_mod, "install_tool_binary", lambda *args, **kwargs: True) + monkeypatch.setattr( + cli_mod, + "configure_selected_tools", + lambda state, tools: {**state, "available_tools": tools}, + ) + monkeypatch.setattr(cli_mod, "validate_all_tools", lambda state: None) + + assert cli_mod.configure_workspace_command() == 0 + assert captured["profile"] == "picked-profile" + def test_multiple_workspaces_configure_all_and_use_first(self, monkeypatch): import ucode.cli as cli_mod diff --git a/tests/test_databricks.py b/tests/test_databricks.py index fffdba5..9936a5a 100644 --- a/tests/test_databricks.py +++ b/tests/test_databricks.py @@ -21,6 +21,7 @@ build_shared_base_urls, build_tool_base_url, ensure_databricks_cli_version, + get_databricks_profiles, get_databricks_token, list_databricks_apps, list_databricks_connections, @@ -405,6 +406,63 @@ def test_error_suggests_logout_when_matching_profile_exists(self, tmp_path, monk assert f"databricks auth login --host {WS} --profile example-profile" in message +class TestGetDatabricksProfiles: + def _patched_run(self, monkeypatch, payload: dict, returncode: int = 0) -> None: + def fake_run(args, **kwargs): + return subprocess.CompletedProcess(args, returncode, stdout=json.dumps(payload)) + + monkeypatch.setattr(db_mod, "run", fake_run) + + def test_keeps_duplicate_hosts_as_separate_entries(self, monkeypatch): + self._patched_run( + monkeypatch, + { + "profiles": [ + {"host": WS, "name": "first", "auth_type": "databricks-cli"}, + {"host": WS, "name": "second", "auth_type": "databricks-cli"}, + { + "host": "https://other.databricks.com", + "name": "third", + "auth_type": "databricks-cli", + }, + ] + }, + ) + profiles = get_databricks_profiles() + assert profiles == [ + (WS, "first"), + (WS, "second"), + ("https://other.databricks.com", "third"), + ] + + def test_skips_pat_profiles(self, monkeypatch): + self._patched_run( + monkeypatch, + { + "profiles": [ + {"host": WS, "name": "oauth", "auth_type": "databricks-cli"}, + {"host": WS, "name": "tokenized", "auth_type": "pat"}, + ] + }, + ) + assert get_databricks_profiles() == [(WS, "oauth")] + + def test_strips_trailing_slash_on_host(self, monkeypatch): + self._patched_run( + monkeypatch, + { + "profiles": [ + {"host": f"{WS}/", "name": "p", "auth_type": "databricks-cli"}, + ] + }, + ) + assert get_databricks_profiles() == [(WS, "p")] + + def test_returns_empty_on_non_zero_exit(self, monkeypatch): + self._patched_run(monkeypatch, {"profiles": []}, returncode=1) + assert get_databricks_profiles() == [] + + class TestListDatabricksConnections: def test_lists_paginated_connections_with_workspace_env(self, monkeypatch): calls: list[dict] = [] diff --git a/tests/test_ui.py b/tests/test_ui.py index fac7838..935455f 100644 --- a/tests/test_ui.py +++ b/tests/test_ui.py @@ -6,7 +6,9 @@ from unittest.mock import patch import pytest +import questionary +from ucode import ui as ui_mod from ucode.ui import ( format_duration, format_token_count, @@ -148,13 +150,99 @@ def test_dash_for_empty_cell(self): assert "-" in result +class _StubQuestion: + def __init__(self, answer): + self._answer = answer + + def ask(self): + return self._answer + + class TestPromptForWorkspace: - """Cover the three things `questionary.select(...).ask()` can return: - a (host, profile) tuple, None (cancel or "Enter a different URL"), - or — in some questionary versions — the choice's title string.""" + """Two concerns combined here: + + 1. Layout — that the picker renders a header separator + one row per + profile (including duplicates of the same host) with ljust-padded + profile names. Driven by capturing the choices list passed to + ``questionary.select``. + 2. Return-type defensiveness — `questionary.select(...).ask()` can hand + back a tuple, None, or (in some versions) the choice title string; + only a tuple should be treated as a selection. + """ PROFILES = [("https://a.databricks.com", "prof-a"), ("https://b.databricks.com", "prof-b")] + # ------------------------------------------------------------------ + # Layout assertions + # ------------------------------------------------------------------ + + def _capture_select(self, monkeypatch, answer): + captured: dict = {} + + def fake_select(message, choices, **kwargs): + captured["message"] = message + captured["choices"] = choices + captured["kwargs"] = kwargs + return _StubQuestion(answer) + + monkeypatch.setattr(questionary, "select", fake_select) + monkeypatch.setattr(ui_mod.questionary, "select", fake_select) + return captured + + def test_shows_header_and_each_profile_row(self, monkeypatch): + profiles = [ + ("https://a.cloud.databricks.com", "alpha"), + ("https://b.cloud.databricks.com", "beta-profile-name"), + ] + captured = self._capture_select(monkeypatch, answer=profiles[0]) + url, profile = prompt_for_workspace("setup", profiles) + + assert (url, profile) == profiles[0] + choices = captured["choices"] + # Header (separator), 2 rows, "Enter a different URL" entry. + assert len(choices) == 4 + assert isinstance(choices[0], questionary.Separator) + header = choices[0].title + assert "Profile Name" in header + assert "Workspace URL" in header + # Profile names ljust-padded to the longest name (17 chars). + name_width = max(len(name) for _, name in profiles) + assert "alpha".ljust(name_width) in choices[1].title + assert profiles[0][0] in choices[1].title + assert "beta-profile-name".ljust(name_width) in choices[2].title + assert profiles[1][0] in choices[2].title + # Final fallback entry still present. + assert choices[3].title == "Enter a different URL" + + def test_keeps_duplicate_hosts_as_separate_rows(self, monkeypatch): + profiles = [ + ("https://shared.cloud.databricks.com", "first"), + ("https://shared.cloud.databricks.com", "second"), + ] + captured = self._capture_select(monkeypatch, answer=profiles[1]) + url, profile = prompt_for_workspace("setup", profiles) + + assert (url, profile) == profiles[1] + # Both rows present — duplicates not collapsed. + choices = captured["choices"] + # Filter to choices whose value is a (host, profile) tuple — drops the + # header separator and the trailing "Enter a different URL" entry. + host_choices = [c for c in choices if isinstance(getattr(c, "value", None), tuple)] + assert [c.value for c in host_choices] == profiles + + def test_returns_normalized_url_with_profile(self, monkeypatch): + # Picker handed back a URL with a trailing slash — normalize_workspace_url + # should strip it before returning. + profiles = [("https://example.cloud.databricks.com/", "p")] + self._capture_select(monkeypatch, answer=profiles[0]) + url, profile = prompt_for_workspace("setup", profiles) + assert url == "https://example.cloud.databricks.com" + assert profile == "p" + + # ------------------------------------------------------------------ + # Return-type defensiveness (regression coverage from PR #104) + # ------------------------------------------------------------------ + def test_returns_selected_profile_tuple(self): with patch("ucode.ui.questionary.select") as mock_select: mock_select.return_value.ask.return_value = ( @@ -192,3 +280,28 @@ def test_no_profiles_goes_straight_to_manual_prompt(self): url, profile = prompt_for_workspace("desc", profiles=None) assert url == "https://example.databricks.com" assert profile is None + + # ------------------------------------------------------------------ + # Long-name display clamping (PR #114 review feedback) + # ------------------------------------------------------------------ + + def test_long_profile_name_is_truncated_in_display_only(self, monkeypatch): + # 60-char name — exceeds the 40-char clamp. The displayed row title + # must be truncated with an ellipsis but the value tuple must carry + # the full untruncated name through to configure_shared_state. + long_name = "x" * 60 + profiles = [("https://a.cloud.databricks.com", long_name)] + captured = self._capture_select(monkeypatch, answer=profiles[0]) + url, profile = prompt_for_workspace("setup", profiles) + + assert (url, profile) == profiles[0] + choices = captured["choices"] + # Header + 1 row + "Enter a different URL". + assert len(choices) == 3 + # Display title is truncated to 40 chars (39 of name + "…"). + row_title = choices[1].title + assert long_name not in row_title + assert "…" in row_title + # Value tuple still carries the full name. + assert choices[1].value == profiles[0] +