Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 8 additions & 4 deletions src/ucode/databricks.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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]] = []
Comment thread
anthonyivn2 marked this conversation as resolved.
pat = 0
for p in profiles:
host = (p.get("host") or "").rstrip("/")
Expand All @@ -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:
Comment thread
anthonyivn2 marked this conversation as resolved.
Expand Down
35 changes: 30 additions & 5 deletions src/ucode/ui.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down
33 changes: 33 additions & 0 deletions tests/test_cli.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
58 changes: 58 additions & 0 deletions tests/test_databricks.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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] = []
Expand Down
119 changes: 116 additions & 3 deletions tests/test_ui.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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 = (
Expand Down Expand Up @@ -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]