diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 6a1662ae..d0789093 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -45,7 +45,6 @@ discover_codex_models, discover_gemini_models, discover_model_services, - ensure_ai_gateway, ensure_databricks_auth, ensure_pat_bearer, find_profile_name_for_host, @@ -57,6 +56,7 @@ list_profile_entries, list_tool_provider_services, normalize_workspace_url, + probe_unity_gateway_capabilities, resolve_pat_token, resolve_provider_launch_model, run_databricks_login, @@ -516,7 +516,7 @@ def configure_shared_state( fable_enabled: bool | None = None, databricks_ai_tools_enabled: bool | None = None, ) -> dict: - """Log into Databricks, enforce AI Gateway v2, fetch model lists, persist state. + """Log into Databricks, verify AI Gateway, fetch model lists, persist state. If tools is provided, only fetch models for those tools. Otherwise fetch all. If force_login is True, always run databricks auth login (used by explicit configure). @@ -635,8 +635,9 @@ def configure_shared_state( state["profile"] = profile with spinner("Verifying Unity AI Gateway..."): token = get_databricks_token(workspace, profile) - ensure_ai_gateway(workspace, token) + model_service_probe = probe_unity_gateway_capabilities(workspace, token) print_success("Unity AI Gateway detected") + print_kv("Model service", model_service_probe.detail) want_claude = ( fetch_all or "claude" in tools or "opencode" in tools or "copilot" in tools or "pi" in tools diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index f875a5c7..33e36c2e 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -51,7 +51,7 @@ WINDOWS_DATABRICKS_INSTALL_URL = ( "https://raw.githubusercontent.com/databricks/setup-cli/main/install.ps1" ) -AI_GATEWAY_V2_DOCS_URL = "https://docs.databricks.com/aws/en/ai-gateway/overview-beta" +AI_GATEWAY_DOCS_URL = "https://docs.databricks.com/aws/en/ai-gateway/overview-beta" ANTHROPIC_MODELS_PATH = "/ai-gateway/anthropic/v1/models" # v1.0.0 is the release that ships `databricks aitools`. MIN_DATABRICKS_CLI_VERSION = (1, 0, 0) @@ -3046,18 +3046,58 @@ def fetch_codex_models(workspace: str, token: str) -> list[str]: return models -def _probe_ai_gateway_v2(workspace: str, token: str) -> tuple[bool, str | None]: +class GatewayProbe(NamedTuple): + reachable: bool + detail: str + resource_available: bool = False + + +def _version_neutral_gateway_detail(detail: str) -> str: + detail = re.sub(r"\bv3\b", "model service", detail, flags=re.IGNORECASE) + return re.sub(r"\bv2\b", "legacy endpoint", detail, flags=re.IGNORECASE) + + +def _gateway_probe_result( + payload: dict | list | None, + reason: str | None, + collection_key: str, + resource_name: str, + empty_hint: str | None = None, +) -> GatewayProbe: + if payload is None: + return GatewayProbe(False, _version_neutral_gateway_detail(reason or "unknown error")) + resources = payload.get(collection_key) if isinstance(payload, dict) else None + if resources: + return GatewayProbe(True, f"reachable, accessible {resource_name} returned", True) + detail = f"reachable, no accessible {resource_name}s returned" + if empty_hint: + detail = f"{detail}; {empty_hint}" + return GatewayProbe(True, detail) + + +def _probe_ai_gateway_v2(workspace: str, token: str) -> GatewayProbe: hostname = workspace_hostname(workspace) url = f"https://{hostname}/api/ai-gateway/v2/endpoints?page_size=1" payload, reason = _http_get_json(url, token) - return payload is not None, reason + return _gateway_probe_result( + payload=payload, + reason=reason, + collection_key="endpoints", + resource_name="endpoint", + ) -def _probe_ai_gateway_v3(workspace: str, token: str) -> tuple[bool, str | None]: +def _probe_ai_gateway_v3(workspace: str, token: str) -> GatewayProbe: hostname = workspace_hostname(workspace) url = f"https://{hostname}/api/2.1/unity-catalog/model-services?page_size=1" payload, reason = _http_get_json(url, token) - return payload is not None, reason + return _gateway_probe_result( + payload=payload, + reason=reason, + collection_key="model_services", + resource_name="model service", + empty_hint="check USE CATALOG on system, and USE SCHEMA and EXECUTE on system.ai", + ) def _raise_ai_gateway_auth_failure(workspace: str, reason: str) -> NoReturn: @@ -3069,57 +3109,65 @@ def _raise_ai_gateway_auth_failure(workspace: str, reason: str) -> NoReturn: ) -def _raise_ai_gateway_v3_permission_failure( - workspace: str, v3_reason: str, v2_reason: str | None +def _raise_model_service_permission_failure( + workspace: str, model_service_reason: str, legacy_endpoint_reason: str ) -> NoReturn: raise RuntimeError( - f"Databricks AI Gateway V3 access could not be verified on {workspace} ({v3_reason}). " - f"The V2 fallback also failed ({v2_reason or 'unknown error'}). The V3 probe requires " - "permission to list Unity Catalog model services. Verify USE CATALOG on `system` and " - "USE SCHEMA on `system.ai`." + "Databricks Unity AI Gateway model service access could not be verified on " + f"{workspace} ({model_service_reason}). The legacy endpoint fallback also failed " + f"({legacy_endpoint_reason}). The model service probe requires permission to list " + "Unity Catalog model services. Verify USE CATALOG on `system`, and USE SCHEMA and " + "EXECUTE on `system.ai`." ) -def _raise_ai_gateway_v2_permission_failure( - workspace: str, v2_reason: str, v3_reason: str | None +def _raise_legacy_endpoint_permission_failure( + workspace: str, legacy_endpoint_reason: str, model_service_reason: str ) -> NoReturn: raise RuntimeError( - f"Databricks AI Gateway V2 access could not be verified on {workspace} ({v2_reason}). " - f"The V3 probe also failed ({v3_reason or 'unknown error'}). Verify the caller's " - "workspace permissions for the AI Gateway V2 endpoints listing." + "Databricks Unity AI Gateway legacy endpoint access could not be verified on " + f"{workspace} ({legacy_endpoint_reason}). The model service probe also failed " + f"({model_service_reason}). Verify the caller's workspace permissions for the legacy " + "endpoints listing." ) -def ensure_ai_gateway(workspace: str, token: str) -> None: - """Pass if either AI Gateway V2 or V3 is available.""" - v3_ok, v3_reason = _probe_ai_gateway_v3(workspace, token) - if v3_ok: - return - if v3_reason and _looks_like_definitive_auth_failure(v3_reason): - _raise_ai_gateway_auth_failure(workspace, v3_reason) - - v2_ok, v2_reason = _probe_ai_gateway_v2(workspace, token) - if v2_ok: - return - if v2_reason and _looks_like_definitive_auth_failure(v2_reason): - _raise_ai_gateway_auth_failure(workspace, v2_reason) - if v3_reason and _looks_like_permission_failure(v3_reason): - _raise_ai_gateway_v3_permission_failure(workspace, v3_reason, v2_reason) - if v2_reason and _looks_like_permission_failure(v2_reason): - _raise_ai_gateway_v2_permission_failure(workspace, v2_reason, v3_reason) +def probe_unity_gateway_capabilities(workspace: str, token: str) -> GatewayProbe: + """Return the model service probe after verifying an available gateway path.""" + model_service_probe = _probe_ai_gateway_v3(workspace, token) + if not model_service_probe.reachable and _looks_like_definitive_auth_failure( + model_service_probe.detail + ): + _raise_ai_gateway_auth_failure(workspace, model_service_probe.detail) + if model_service_probe.resource_available: + return model_service_probe + + legacy_endpoint_probe = _probe_ai_gateway_v2(workspace, token) + if legacy_endpoint_probe.reachable: + return model_service_probe + if _looks_like_definitive_auth_failure(legacy_endpoint_probe.detail): + _raise_ai_gateway_auth_failure(workspace, legacy_endpoint_probe.detail) + if _looks_like_permission_failure(model_service_probe.detail): + _raise_model_service_permission_failure( + workspace, model_service_probe.detail, legacy_endpoint_probe.detail + ) + if _looks_like_permission_failure(legacy_endpoint_probe.detail): + _raise_legacy_endpoint_permission_failure( + workspace, legacy_endpoint_probe.detail, model_service_probe.detail + ) raise RuntimeError( - "Databricks AI Gateway is not enabled on this workspace: neither V3 " - f"({v3_reason or 'unknown error'}) nor V2 ({v2_reason or 'unknown error'}) is available. " - f"See {AI_GATEWAY_V2_DOCS_URL}" + "Databricks Unity AI Gateway is not enabled on this workspace: neither model services " + f"({model_service_probe.detail}) nor legacy endpoints ({legacy_endpoint_probe.detail}) " + f"are available. See {AI_GATEWAY_DOCS_URL}" ) def _looks_like_definitive_auth_failure(reason: str) -> bool: """True when retrying another workspace API cannot rescue this token. - A 403 can be endpoint-specific authorization, so the version-agnostic - preflight must still try V3 before surfacing it as an auth failure. + A 403 can be endpoint-specific authorization, so the preflight must still + try the fallback before surfacing it as an auth failure. """ if "HTTP 401" in reason: return True diff --git a/tests/test_cli.py b/tests/test_cli.py index 8bb1ea9b..a59f3015 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -13,7 +13,9 @@ import pytest from typer.testing import CliRunner +import ucode.databricks as db_mod from ucode.cli import app +from ucode.databricks import GatewayProbe _ANSI_RE = re.compile(r"\x1b\[[0-9;]*m") @@ -29,6 +31,8 @@ def _strip_ansi(text: str) -> str: TOOLS = ["codex", "claude", "gemini", "opencode"] +MODEL_SERVICE_PROBE = GatewayProbe(True, "reachable, accessible model service returned", True) + def _jwt(expires_at: float) -> str: payload = base64.urlsafe_b64encode(json.dumps({"exp": expires_at}).encode()).decode() @@ -2466,7 +2470,9 @@ def _stub_deps(monkeypatch, *, pat_token, existing_state=None): monkeypatch.setattr(cli_mod, "resolve_pat_token", lambda p: pat_token) monkeypatch.setattr(cli_mod, "find_profile_name_for_host", lambda w: None) monkeypatch.setattr(cli_mod, "get_databricks_token", lambda w, p: "token") - monkeypatch.setattr(cli_mod, "ensure_ai_gateway", lambda w, t: None) + monkeypatch.setattr( + cli_mod, "probe_unity_gateway_capabilities", lambda w, t: MODEL_SERVICE_PROBE + ) monkeypatch.setattr(cli_mod, "discover_model_services", lambda w, t: ({}, [], [], [], None)) monkeypatch.setattr(cli_mod, "discover_claude_models", lambda w, t: ({}, None)) monkeypatch.setattr(cli_mod, "discover_gemini_models", lambda w, t: ([], None)) @@ -2489,6 +2495,112 @@ def test_use_pat_exports_bearer_and_skips_login(self, monkeypatch): assert state["use_pat"] is True assert saved and saved[-1]["use_pat"] is True + def test_prints_model_service_and_omits_unneeded_legacy_probe(self, monkeypatch, capsys): + cli_mod, *_ = self._stub_deps(monkeypatch, pat_token="dapi-pat") + + cli_mod.configure_shared_state(self.WS, profile="DEFAULT") + + output = _strip_ansi(capsys.readouterr().out) + assert "Model service: reachable, accessible model service returned" in output + assert "(Legacy) endpoints:" not in output + + @pytest.mark.parametrize( + ("responses", "expected_model_service"), + [ + ( + [({"model_services": [{"name": "model-services/system.ai.gpt-5"}]}, None)], + "reachable, accessible model service returned", + ), + ( + [({}, None), ({"endpoints": []}, None)], + "reachable, no accessible model services returned; check USE CATALOG on system, " + "and USE SCHEMA and EXECUTE on system.ai", + ), + ( + [ + (None, "HTTP 403 Forbidden"), + ({"endpoints": [{"name": "databricks-gpt-5"}]}, None), + ], + "HTTP 403 Forbidden", + ), + ], + ids=[ + "model-service-resource", + "model-service-empty-legacy-empty", + "model-service-forbidden-legacy-resource", + ], + ) + def test_prints_local_gateway_probe_scenarios( + self, + monkeypatch, + capsys, + responses, + expected_model_service, + ): + cli_mod, *_ = self._stub_deps(monkeypatch, pat_token="dapi-pat") + response_iter = iter(responses) + monkeypatch.setattr(db_mod, "_http_get_json", lambda url, token: next(response_iter)) + monkeypatch.setattr( + cli_mod, "probe_unity_gateway_capabilities", db_mod.probe_unity_gateway_capabilities + ) + + cli_mod.configure_shared_state(self.WS, profile="DEFAULT") + + output = " ".join(_strip_ansi(capsys.readouterr().out).split()) + assert "Unity AI Gateway detected" in output + assert f"Model service: {expected_model_service}" in output + assert "(Legacy) endpoints:" not in output + assert "V2" not in output + assert "V3" not in output + + @pytest.mark.parametrize( + ("responses", "error_match"), + [ + ( + [ + ({}, None), + ( + None, + "HTTP 404 Not Found: AI Gateway V2 is not available for CSP-enabled " + "workspaces", + ), + ], + "no accessible model services", + ), + ( + [ + (None, "HTTP 404 Not Found: V3 unavailable"), + (None, "HTTP 404 Not Found: V2 unavailable"), + ], + "neither model services", + ), + ([(None, "HTTP 401 Unauthorized")], "rejected the access token"), + ], + ids=[ + "model-service-empty-legacy-unavailable", + "neither-path-reachable", + "invalid-token", + ], + ) + def test_local_gateway_probe_failures_do_not_print_success( + self, monkeypatch, capsys, responses, error_match + ): + cli_mod, *_ = self._stub_deps(monkeypatch, pat_token="dapi-pat") + response_iter = iter(responses) + monkeypatch.setattr(db_mod, "_http_get_json", lambda url, token: next(response_iter)) + monkeypatch.setattr( + cli_mod, "probe_unity_gateway_capabilities", db_mod.probe_unity_gateway_capabilities + ) + + with pytest.raises(RuntimeError, match=error_match) as excinfo: + cli_mod.configure_shared_state(self.WS, profile="DEFAULT") + + output = _strip_ansi(capsys.readouterr().out) + assert "Unity AI Gateway detected" not in output + message = str(excinfo.value) + assert "v2" not in message.lower() + assert "v3" not in message.lower() + def test_use_pat_without_pat_profile_raises(self, monkeypatch): cli_mod, logins, _, _ = self._stub_deps(monkeypatch, pat_token=None) @@ -2787,7 +2899,9 @@ def _stub_external_deps(monkeypatch): monkeypatch.setattr(cli_mod, "ensure_databricks_auth", lambda w, p=None: None) monkeypatch.setattr(cli_mod, "find_profile_name_for_host", lambda w: None) monkeypatch.setattr(cli_mod, "get_databricks_token", lambda w, p: "token") - monkeypatch.setattr(cli_mod, "ensure_ai_gateway", lambda w, t: None) + monkeypatch.setattr( + cli_mod, "probe_unity_gateway_capabilities", lambda w, t: MODEL_SERVICE_PROBE + ) monkeypatch.setattr(cli_mod, "discover_model_services", lambda w, t: ({}, [], [], [], None)) monkeypatch.setattr(cli_mod, "discover_claude_models", lambda w, t: ({}, None)) monkeypatch.setattr(cli_mod, "discover_gemini_models", lambda w, t: ([], None)) @@ -2847,7 +2961,9 @@ def _stub(monkeypatch): monkeypatch.setattr(cli_mod, "run_databricks_login", lambda w, p: None) monkeypatch.setattr(cli_mod, "find_profile_name_for_host", lambda w: None) monkeypatch.setattr(cli_mod, "get_databricks_token", lambda w, p: "token") - monkeypatch.setattr(cli_mod, "ensure_ai_gateway", lambda w, t: None) + monkeypatch.setattr( + cli_mod, "probe_unity_gateway_capabilities", lambda w, t: MODEL_SERVICE_PROBE + ) monkeypatch.setattr(cli_mod, "build_shared_base_urls", lambda w: {}) monkeypatch.setattr(cli_mod, "save_state", lambda s: None) @@ -2905,7 +3021,11 @@ def _f(*a, **k): monkeypatch.setattr(cli_mod, "run_databricks_login", _boom("run_databricks_login")) monkeypatch.setattr(cli_mod, "ensure_pat_bearer", _boom("ensure_pat_bearer")) monkeypatch.setattr(cli_mod, "get_databricks_token", _boom("get_databricks_token")) - monkeypatch.setattr(cli_mod, "ensure_ai_gateway", _boom("ensure_ai_gateway")) + monkeypatch.setattr( + cli_mod, + "probe_unity_gateway_capabilities", + _boom("probe_unity_gateway_capabilities"), + ) monkeypatch.setattr(cli_mod, "discover_model_services", _boom("discover_model_services")) monkeypatch.setattr(cli_mod, "discover_codex_models", _boom("discover_codex_models")) monkeypatch.setattr(cli_mod, "find_profile_name_for_host", lambda w: "resolved") diff --git a/tests/test_databricks.py b/tests/test_databricks.py index 2d9f61a5..0df389cc 100644 --- a/tests/test_databricks.py +++ b/tests/test_databricks.py @@ -1915,21 +1915,41 @@ def fake_run(args, **kwargs): list_databricks_apps(WS) -class TestEnsureAiGateway: - def test_v3_only_workspace_succeeds_without_v2_probe(self, monkeypatch): +class TestProbeUnityGatewayCapabilities: + def test_model_service_resource_skips_legacy_probe(self, monkeypatch): calls: list[str] = [] def fake_get(url, token): calls.append(url) - return {"model_services": []}, None + return {"model_services": [{"name": "model-services/system.ai.gpt-5"}]}, None monkeypatch.setattr(db_mod, "_http_get_json", fake_get) - db_mod.ensure_ai_gateway(WS, "fake-token") + model_service_probe = db_mod.probe_unity_gateway_capabilities(WS, "fake-token") - assert calls == [f"https://{WS_HOST}/api/2.1/unity-catalog/model-services?page_size=1"] + assert model_service_probe == db_mod.GatewayProbe( + True, "reachable, accessible model service returned", True + ) + assert calls == [ + f"https://{WS_HOST}/api/2.1/unity-catalog/model-services?page_size=1", + ] + + def test_empty_model_service_response_includes_permission_hint(self, monkeypatch): + monkeypatch.setattr( + db_mod, + "_http_get_json", + lambda url, token: ({}, None), + ) + + model_service_probe = db_mod.probe_unity_gateway_capabilities(WS, "fake-token") - def test_v2_only_workspace_succeeds_after_v3_probe(self, monkeypatch): + assert model_service_probe == db_mod.GatewayProbe( + True, + "reachable, no accessible model services returned; check USE CATALOG on system, and " + "USE SCHEMA and EXECUTE on system.ai", + ) + + def test_legacy_only_workspace_returns_model_service_probe(self, monkeypatch): calls: list[str] = [] def fake_get(url, token): @@ -1940,14 +1960,15 @@ def fake_get(url, token): monkeypatch.setattr(db_mod, "_http_get_json", fake_get) - db_mod.ensure_ai_gateway(WS, "fake-token") + model_service_probe = db_mod.probe_unity_gateway_capabilities(WS, "fake-token") + assert model_service_probe == db_mod.GatewayProbe(False, "HTTP 404: Not Found") assert calls == [ f"https://{WS_HOST}/api/2.1/unity-catalog/model-services?page_size=1", f"https://{WS_HOST}/api/ai-gateway/v2/endpoints?page_size=1", ] - def test_v3_forbidden_still_succeeds_when_v2_is_available(self, monkeypatch): + def test_model_service_forbidden_still_succeeds_when_legacy_is_available(self, monkeypatch): calls: list[str] = [] def fake_get(url, token): @@ -1958,29 +1979,58 @@ def fake_get(url, token): monkeypatch.setattr(db_mod, "_http_get_json", fake_get) - db_mod.ensure_ai_gateway(WS, "fake-token") + model_service_probe = db_mod.probe_unity_gateway_capabilities(WS, "fake-token") + assert model_service_probe == db_mod.GatewayProbe(False, "HTTP 403: Forbidden") assert calls == [ f"https://{WS_HOST}/api/2.1/unity-catalog/model-services?page_size=1", f"https://{WS_HOST}/api/ai-gateway/v2/endpoints?page_size=1", ] + def test_empty_model_service_requires_reachable_legacy_fallback(self, monkeypatch): + responses = iter( + [ + ({}, None), + (None, "HTTP 404: AI Gateway V2 is not available for CSP-enabled workspaces"), + ] + ) + monkeypatch.setattr( + db_mod, + "_http_get_json", + lambda url, token: next(responses), + ) + + with pytest.raises(RuntimeError, match="no accessible model services") as excinfo: + db_mod.probe_unity_gateway_capabilities(WS, "fake-token") + + message = str(excinfo.value) + assert "HTTP 404: AI Gateway legacy endpoint is not available" in message + assert "v2" not in message.lower() + assert "v3" not in message.lower() + def test_neither_gateway_available_raises(self, monkeypatch): - reasons = iter(["HTTP 404: V3 missing", "HTTP 404: V2 missing"]) + reasons = iter( + [ + "HTTP 404: V3 unavailable", + "HTTP 404: V2 unavailable", + ] + ) monkeypatch.setattr( db_mod, "_http_get_json", lambda url, token: (None, next(reasons)), ) - with pytest.raises(RuntimeError, match="neither V3") as excinfo: - db_mod.ensure_ai_gateway(WS, "fake-token") + with pytest.raises(RuntimeError, match="neither model services") as excinfo: + db_mod.probe_unity_gateway_capabilities(WS, "fake-token") message = str(excinfo.value) - assert "HTTP 404: V2 missing" in message - assert "HTTP 404: V3 missing" in message + assert "HTTP 404: model service unavailable" in message + assert "HTTP 404: legacy endpoint unavailable" in message + assert "v2" not in message.lower() + assert "v3" not in message.lower() - def test_v3_auth_failure_does_not_probe_v2(self, monkeypatch): + def test_model_service_auth_failure_does_not_probe_legacy(self, monkeypatch): calls: list[str] = [] def fake_get(url, token): @@ -1990,12 +2040,19 @@ def fake_get(url, token): monkeypatch.setattr(db_mod, "_http_get_json", fake_get) with pytest.raises(RuntimeError, match="rejected"): - db_mod.ensure_ai_gateway(WS, "fake-token") + db_mod.probe_unity_gateway_capabilities(WS, "fake-token") assert calls == [f"https://{WS_HOST}/api/2.1/unity-catalog/model-services?page_size=1"] - def test_v3_forbidden_and_v2_unavailable_reports_permission_error(self, monkeypatch): - reasons = iter(["HTTP 403: Missing Unity Catalog grants", "HTTP 404: V2 missing"]) + def test_model_service_forbidden_and_legacy_unavailable_reports_permission_error( + self, monkeypatch + ): + reasons = iter( + [ + "HTTP 403: Missing Unity Catalog grants", + "HTTP 404: legacy endpoints unavailable", + ] + ) monkeypatch.setattr( db_mod, "_http_get_json", @@ -2003,15 +2060,25 @@ def test_v3_forbidden_and_v2_unavailable_reports_permission_error(self, monkeypa ) with pytest.raises(RuntimeError, match="permission") as excinfo: - db_mod.ensure_ai_gateway(WS, "fake-token") + db_mod.probe_unity_gateway_capabilities(WS, "fake-token") message = str(excinfo.value) assert "USE SCHEMA" in message + assert "EXECUTE" in message + assert "v2" not in message.lower() + assert "v3" not in message.lower() assert "rejected the access token" not in message assert "not enabled" not in message - def test_v2_forbidden_and_v3_unavailable_reports_permission_error(self, monkeypatch): - reasons = iter(["HTTP 404: V3 missing", "HTTP 403: V2 forbidden"]) + def test_legacy_forbidden_and_model_service_unavailable_reports_permission_error( + self, monkeypatch + ): + reasons = iter( + [ + "HTTP 404: V3 unavailable", + "HTTP 403: V2 forbidden", + ] + ) monkeypatch.setattr( db_mod, "_http_get_json", @@ -2019,16 +2086,18 @@ def test_v2_forbidden_and_v3_unavailable_reports_permission_error(self, monkeypa ) with pytest.raises(RuntimeError, match="workspace permissions") as excinfo: - db_mod.ensure_ai_gateway(WS, "fake-token") + db_mod.probe_unity_gateway_capabilities(WS, "fake-token") message = str(excinfo.value) - assert "V2 access could not be verified" in message + assert "legacy endpoint access could not be verified" in message + assert "v2" not in message.lower() + assert "v3" not in message.lower() assert "USE SCHEMA" not in message class TestHttpGetJsonReason: """The `reason` string returned by `_http_get_json` must include the response body - so callers (e.g. ensure_ai_gateway) can route on it. Before issue #84's fix + so callers (e.g. the Unity Gateway capability probe) can route on it. Before issue #84's fix the body was logged only when UCODE_DEBUG=1 and dropped from the bubbled error.""" @staticmethod diff --git a/tests/test_e2e.py b/tests/test_e2e.py index 3b9319eb..a578e5b2 100644 --- a/tests/test_e2e.py +++ b/tests/test_e2e.py @@ -27,7 +27,6 @@ build_tool_base_url, discover_model_services, discover_sql_warehouses, - ensure_ai_gateway, fetch_ai_gateway_claude_models, fetch_codex_models, fetch_gemini_models, @@ -35,6 +34,7 @@ is_model_provider_feature_unavailable, list_model_provider_services, list_tool_provider_services, + probe_unity_gateway_capabilities, service_usable_for_tool, workspace_hostname, ) @@ -137,8 +137,8 @@ def test_get_token_returns_non_empty_string(self, e2e_token): class TestAiGateway: - def test_ensure_ai_gateway_does_not_raise(self, e2e_workspace, e2e_token): - ensure_ai_gateway(e2e_workspace, e2e_token) + def test_probe_unity_gateway_capabilities_does_not_raise(self, e2e_workspace, e2e_token): + probe_unity_gateway_capabilities(e2e_workspace, e2e_token) def test_workspace_hostname_resolves(self, e2e_workspace): hostname = workspace_hostname(e2e_workspace) diff --git a/tests/test_managed_wizard.py b/tests/test_managed_wizard.py index 013cdc9d..0fcbebac 100644 --- a/tests/test_managed_wizard.py +++ b/tests/test_managed_wizard.py @@ -2857,8 +2857,11 @@ def test_publish_file_flag_is_forwarded(self): assert publish.call_args.kwargs["file_path"] == "/tmp/cfg.json" def test_publish_error_exits_nonzero_with_a_message(self): - with patch.object( - cli_mod, "publish_command", side_effect=RuntimeError("no config authored") + with ( + patch("ucode.cli.install_databricks_cli"), + patch.object( + cli_mod, "publish_command", side_effect=RuntimeError("no config authored") + ), ): result = runner.invoke(app, ["publish"]) assert result.exit_code == 1 @@ -2866,7 +2869,10 @@ def test_publish_error_exits_nonzero_with_a_message(self): def test_successful_publish_exits_zero(self): # Same trap as `setup`: `typer.Exit` subclasses RuntimeError, so raising it inside the # command's try block would report success as "ERROR 0". - with patch.object(cli_mod, "publish_command", return_value=0): + with ( + patch("ucode.cli.install_databricks_cli"), + patch.object(cli_mod, "publish_command", return_value=0), + ): result = runner.invoke(app, ["publish"]) assert result.exit_code == 0 assert "ERROR" not in result.output