chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,412 @@
|
||||
"""Tests for omnigent.runtime.credentials.databricks."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
from omnigent.runtime.credentials.databricks import (
|
||||
WorkspaceCreds,
|
||||
resolve_databricks_workspace,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _clean_env(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
"""
|
||||
Strip every credential-related env var before each test so the
|
||||
resolver cannot accidentally pick up the developer's real
|
||||
credentials, and short-circuit the SDK path so tests don't
|
||||
perform real OAuth / network authentication.
|
||||
|
||||
Why the SDK short-circuit: ``databricks.sdk.config.Config(...)
|
||||
.authenticate()`` walks every supported auth_type — PAT, OAuth-
|
||||
U2M, Azure CLI, OIDC, IMDS, etc. Several of those make
|
||||
network calls or shell out to external CLIs even when the
|
||||
resolver intends to fall through to the cfg-file path. Without
|
||||
this monkeypatch the suite takes ~35 minutes; with it, ~0.4s.
|
||||
Each test that wants to exercise the SDK path explicitly
|
||||
overrides this monkeypatch (see ``test_resolves_via_sdk_*``).
|
||||
"""
|
||||
for var in (
|
||||
"DATABRICKS_CONFIG_FILE",
|
||||
"DATABRICKS_CONFIG_PROFILE",
|
||||
):
|
||||
monkeypatch.delenv(var, raising=False)
|
||||
|
||||
def _raise_value_error(*_args: object, **_kwargs: object) -> None:
|
||||
raise ValueError("SDK path disabled in tests by default")
|
||||
|
||||
monkeypatch.setattr("databricks.sdk.config.Config", _raise_value_error)
|
||||
|
||||
|
||||
def _write_cfg(tmp_path: Path, body: str) -> Path:
|
||||
"""
|
||||
Write *body* to a temporary ``.databrickscfg`` and return its path.
|
||||
|
||||
:param tmp_path: pytest's per-test temp directory.
|
||||
:param body: The exact INI text to write.
|
||||
:returns: The :class:`Path` to the written config file.
|
||||
"""
|
||||
cfg = tmp_path / "databrickscfg"
|
||||
cfg.write_text(body)
|
||||
return cfg
|
||||
|
||||
|
||||
def test_openai_env_vars_are_ignored(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None:
|
||||
# OPENAI_BASE_URL / OPENAI_API_KEY are used elsewhere in the
|
||||
# codebase to point at full serving-endpoints URLs; this resolver
|
||||
# MUST NOT treat them as workspace creds. If it did, a downstream
|
||||
# caller appending the gateway path would produce a malformed URL
|
||||
# like .../serving-endpoints/ai-gateway/mlflow/v1.
|
||||
monkeypatch.setenv("OPENAI_BASE_URL", "https://example.com/serving-endpoints")
|
||||
monkeypatch.setenv("OPENAI_API_KEY", "openai-token")
|
||||
cfg = _write_cfg(
|
||||
tmp_path,
|
||||
"[DEFAULT]\nhost = https://default.example.com\ntoken = default-token\n",
|
||||
)
|
||||
monkeypatch.setenv("DATABRICKS_CONFIG_FILE", str(cfg))
|
||||
|
||||
creds = resolve_databricks_workspace(profile=None)
|
||||
|
||||
# Resolver ignored the OPENAI_* env vars and went straight to the
|
||||
# cfg file. If this returned the OpenAI host instead, the resolver
|
||||
# is silently mis-using OpenAI env vars as workspace credentials.
|
||||
assert creds.host == "https://default.example.com"
|
||||
assert creds.token == "default-token"
|
||||
|
||||
|
||||
def test_resolves_named_profile_from_cfg(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None:
|
||||
cfg = _write_cfg(
|
||||
tmp_path,
|
||||
("[dev]\nhost = https://dev.example.com\ntoken = dev-token\n"),
|
||||
)
|
||||
monkeypatch.setenv("DATABRICKS_CONFIG_FILE", str(cfg))
|
||||
|
||||
creds = resolve_databricks_workspace(profile="dev")
|
||||
|
||||
# Named-profile branch returned the [dev] section's exact values.
|
||||
assert creds == WorkspaceCreds(host="https://dev.example.com", token="dev-token")
|
||||
|
||||
|
||||
def test_resolves_default_when_profile_is_none(
|
||||
monkeypatch: pytest.MonkeyPatch, tmp_path: Path
|
||||
) -> None:
|
||||
cfg = _write_cfg(
|
||||
tmp_path,
|
||||
("[DEFAULT]\nhost = https://default.example.com\ntoken = default-token\n"),
|
||||
)
|
||||
monkeypatch.setenv("DATABRICKS_CONFIG_FILE", str(cfg))
|
||||
|
||||
creds = resolve_databricks_workspace(profile=None)
|
||||
|
||||
# profile=None went straight to [DEFAULT] and pulled both values.
|
||||
assert creds == WorkspaceCreds(host="https://default.example.com", token="default-token")
|
||||
|
||||
|
||||
def test_named_profile_overrides_default(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None:
|
||||
cfg = _write_cfg(
|
||||
tmp_path,
|
||||
(
|
||||
"[DEFAULT]\n"
|
||||
"host = https://default.example.com\n"
|
||||
"token = default-token\n"
|
||||
"\n"
|
||||
"[dev]\n"
|
||||
"host = https://dev.example.com\n"
|
||||
"token = dev-token\n"
|
||||
),
|
||||
)
|
||||
monkeypatch.setenv("DATABRICKS_CONFIG_FILE", str(cfg))
|
||||
|
||||
creds = resolve_databricks_workspace(profile="dev")
|
||||
|
||||
# When the named profile exists AND has both fields, it wins over
|
||||
# [DEFAULT]. If this returned the DEFAULT host instead, the
|
||||
# resolver is silently ignoring the requested profile.
|
||||
assert creds.host == "https://dev.example.com"
|
||||
assert creds.token == "dev-token"
|
||||
|
||||
|
||||
def test_absent_named_profile_raises_not_falls_back_to_default(
|
||||
monkeypatch: pytest.MonkeyPatch, tmp_path: Path
|
||||
) -> None:
|
||||
# [DEFAULT] has valid creds but [ghost] doesn't exist. The resolver
|
||||
# must raise OSError rather than silently routing to DEFAULT — a
|
||||
# typo in --profile would otherwise send requests to a completely
|
||||
# different workspace without any error.
|
||||
cfg = _write_cfg(
|
||||
tmp_path,
|
||||
("[DEFAULT]\nhost = https://default.example.com\ntoken = default-token\n"),
|
||||
)
|
||||
monkeypatch.setenv("DATABRICKS_CONFIG_FILE", str(cfg))
|
||||
|
||||
with pytest.raises(OSError) as excinfo:
|
||||
resolve_databricks_workspace(profile="ghost")
|
||||
|
||||
msg = str(excinfo.value)
|
||||
assert "ghost" in msg
|
||||
assert str(cfg) in msg
|
||||
|
||||
|
||||
def test_databricks_config_profile_env_var_typo_raises_not_falls_back_to_default(
|
||||
monkeypatch: pytest.MonkeyPatch, tmp_path: Path
|
||||
) -> None:
|
||||
# Simulates: `omnigent run foo.yaml --profile typo-profile --model databricks/m`
|
||||
# _propagate_profile_to_environment sets DATABRICKS_CONFIG_PROFILE="typo-profile".
|
||||
# DatabricksAdapter calls resolve_databricks_workspace(None).
|
||||
# Without the effective_profile fix, the configparser path receives profile=None,
|
||||
# skips the named-section check, and silently returns [DEFAULT] creds.
|
||||
cfg = _write_cfg(
|
||||
tmp_path,
|
||||
"[DEFAULT]\nhost = https://default.example.com\ntoken = default-token\n",
|
||||
)
|
||||
monkeypatch.setenv("DATABRICKS_CONFIG_FILE", str(cfg))
|
||||
monkeypatch.setenv("DATABRICKS_CONFIG_PROFILE", "typo-profile")
|
||||
|
||||
with pytest.raises(OSError) as excinfo:
|
||||
resolve_databricks_workspace(profile=None)
|
||||
|
||||
msg = str(excinfo.value)
|
||||
assert "typo-profile" in msg
|
||||
assert str(cfg) in msg
|
||||
|
||||
|
||||
def test_strips_trailing_slash_from_host(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None:
|
||||
cfg = _write_cfg(
|
||||
tmp_path,
|
||||
("[dev]\nhost = https://dev.example.com///\ntoken = dev-token\n"),
|
||||
)
|
||||
monkeypatch.setenv("DATABRICKS_CONFIG_FILE", str(cfg))
|
||||
|
||||
creds = resolve_databricks_workspace(profile="dev")
|
||||
|
||||
# Three trailing slashes in the cfg, zero in the result. Callers
|
||||
# rely on this normalization to use the host directly as an
|
||||
# OpenAI base_url without further cleanup.
|
||||
assert creds.host == "https://dev.example.com"
|
||||
|
||||
|
||||
def test_raises_when_all_sources_empty(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None:
|
||||
# Point DATABRICKS_CONFIG_FILE at a path that does not exist so
|
||||
# the resolver cannot accidentally pick up the developer's real
|
||||
# ~/.databrickscfg. The autouse fixture has already disabled the
|
||||
# SDK path.
|
||||
missing_cfg = tmp_path / "does-not-exist"
|
||||
monkeypatch.setenv("DATABRICKS_CONFIG_FILE", str(missing_cfg))
|
||||
|
||||
with pytest.raises(OSError) as excinfo:
|
||||
resolve_databricks_workspace(profile="dev")
|
||||
|
||||
msg = str(excinfo.value)
|
||||
# The error must name every source that was checked so the
|
||||
# caller can debug exactly which piece is missing.
|
||||
assert "databricks-sdk" in msg
|
||||
assert "[dev]" in msg or "profile [dev]" in msg
|
||||
assert "[DEFAULT]" in msg
|
||||
assert str(missing_cfg) in msg
|
||||
|
||||
|
||||
def test_raises_when_cfg_section_missing_token(
|
||||
monkeypatch: pytest.MonkeyPatch, tmp_path: Path
|
||||
) -> None:
|
||||
cfg = _write_cfg(
|
||||
tmp_path,
|
||||
(
|
||||
"[dev]\nhost = https://dev.example.com\n"
|
||||
# NOTE: no token line
|
||||
),
|
||||
)
|
||||
monkeypatch.setenv("DATABRICKS_CONFIG_FILE", str(cfg))
|
||||
|
||||
# The [dev] section exists but is missing the token. The resolver
|
||||
# must NOT silently substitute an empty token — it must treat the
|
||||
# section as unresolved and (since [DEFAULT] is also missing)
|
||||
# raise OSError.
|
||||
with pytest.raises(OSError):
|
||||
resolve_databricks_workspace(profile="dev")
|
||||
|
||||
|
||||
def test_malformed_named_section_does_not_silently_fall_back_to_default(
|
||||
monkeypatch: pytest.MonkeyPatch, tmp_path: Path
|
||||
) -> None:
|
||||
# Two sections exist: a malformed [dev] (missing token) and a
|
||||
# complete [DEFAULT]. The resolver must NOT silently use DEFAULT
|
||||
# when the user explicitly asked for [dev] — that would send the
|
||||
# caller to a different workspace than they requested, which is
|
||||
# very hard to debug because everything "works" but talks to
|
||||
# the wrong place.
|
||||
cfg = _write_cfg(
|
||||
tmp_path,
|
||||
(
|
||||
"[DEFAULT]\n"
|
||||
"host = https://default.example.com\n"
|
||||
"token = default-token\n"
|
||||
"\n"
|
||||
"[dev]\n"
|
||||
"host = https://dev.example.com\n"
|
||||
# NOTE: no token line — section present but invalid
|
||||
),
|
||||
)
|
||||
monkeypatch.setenv("DATABRICKS_CONFIG_FILE", str(cfg))
|
||||
|
||||
with pytest.raises(OSError) as excinfo:
|
||||
resolve_databricks_workspace(profile="dev")
|
||||
|
||||
msg = str(excinfo.value)
|
||||
# Error must name the offending profile so the user knows which
|
||||
# section to fix.
|
||||
assert "[dev]" in msg
|
||||
# Error should mention 'malformed' or what's missing so the user
|
||||
# can debug fast.
|
||||
assert "malformed" in msg.lower() or "token" in msg
|
||||
|
||||
|
||||
def test_resolves_via_sdk_when_sdk_returns_creds(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
# Override the autouse fixture's SDK short-circuit so this test
|
||||
# exercises the SDK branch end-to-end (with a fake Config that
|
||||
# returns plausible OAuth-style results — proves the resolver
|
||||
# works for ``auth_type = databricks-cli`` profiles whose cfg
|
||||
# sections have NO static ``token`` field).
|
||||
class _FakeConfig:
|
||||
def __init__(self, *, profile: str | None) -> None:
|
||||
self.host = "https://sdk.example.com/" # trailing slash on purpose
|
||||
|
||||
def authenticate(self) -> dict[str, str]:
|
||||
return {"Authorization": "Bearer sdk-minted-token"}
|
||||
|
||||
monkeypatch.setattr("databricks.sdk.config.Config", _FakeConfig)
|
||||
|
||||
creds = resolve_databricks_workspace(profile="some-oauth-profile")
|
||||
|
||||
# SDK path returned a freshly-minted bearer; trailing slash on
|
||||
# the host was stripped on the way out.
|
||||
assert creds == WorkspaceCreds(host="https://sdk.example.com", token="sdk-minted-token")
|
||||
|
||||
|
||||
def test_sdk_failure_falls_through_to_cfg(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None:
|
||||
# Autouse fixture already sets Config to raise ValueError. Provide
|
||||
# a cfg file with a plain-PAT [dev] section. The resolver must
|
||||
# catch the SDK ValueError and fall through to the configparser
|
||||
# path rather than propagating the failure.
|
||||
cfg = _write_cfg(
|
||||
tmp_path,
|
||||
"[dev]\nhost = https://cfg-fallback.example.com\ntoken = cfg-token\n",
|
||||
)
|
||||
monkeypatch.setenv("DATABRICKS_CONFIG_FILE", str(cfg))
|
||||
|
||||
creds = resolve_databricks_workspace(profile="dev")
|
||||
|
||||
# If this returned None or raised, the resolver isn't catching
|
||||
# the SDK's ValueError properly — the cfg-file fallback would
|
||||
# never run for any setup where the SDK initialization fails.
|
||||
assert creds.host == "https://cfg-fallback.example.com"
|
||||
assert creds.token == "cfg-token"
|
||||
|
||||
|
||||
def test_sdk_non_bearer_auth_falls_through(
|
||||
monkeypatch: pytest.MonkeyPatch, tmp_path: Path
|
||||
) -> None:
|
||||
# Some auth schemes (Basic, etc.) return non-Bearer Authorization
|
||||
# headers. The resolver doesn't support those — it must fall
|
||||
# through to the cfg-file path rather than returning a malformed
|
||||
# token.
|
||||
class _NonBearerConfig:
|
||||
def __init__(self, *, profile: str | None) -> None:
|
||||
self.host = "https://sdk.example.com"
|
||||
|
||||
def authenticate(self) -> dict[str, str]:
|
||||
return {"Authorization": "Basic some-base64-blob"}
|
||||
|
||||
monkeypatch.setattr("databricks.sdk.config.Config", _NonBearerConfig)
|
||||
cfg = _write_cfg(
|
||||
tmp_path,
|
||||
"[dev]\nhost = https://cfg.example.com\ntoken = cfg-token\n",
|
||||
)
|
||||
monkeypatch.setenv("DATABRICKS_CONFIG_FILE", str(cfg))
|
||||
|
||||
creds = resolve_databricks_workspace(profile="dev")
|
||||
|
||||
# SDK returned Basic auth (unsupported) → resolver fell through
|
||||
# to the cfg path. If this returned an "sdk.example.com" host,
|
||||
# the resolver is silently accepting non-Bearer schemes.
|
||||
assert creds.host == "https://cfg.example.com"
|
||||
|
||||
|
||||
def test_sdk_value_error_does_not_emit_warning(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
tmp_path: Path,
|
||||
caplog: pytest.LogCaptureFixture,
|
||||
) -> None:
|
||||
# Regression: SDK ValueError used to log at WARNING with
|
||||
# exc_info=True, dumping a 40-line traceback ahead of the clean
|
||||
# ClickException. INFO keeps it in cli-*.log but off stderr.
|
||||
cfg = _write_cfg(
|
||||
tmp_path,
|
||||
"[dev]\nhost = https://cfg.example.com\ntoken = cfg-token\n",
|
||||
)
|
||||
monkeypatch.setenv("DATABRICKS_CONFIG_FILE", str(cfg))
|
||||
|
||||
with caplog.at_level(logging.DEBUG, logger="omnigent.runtime.credentials.databricks"):
|
||||
resolve_databricks_workspace(profile="dev")
|
||||
|
||||
module_records = [
|
||||
r for r in caplog.records if r.name == "omnigent.runtime.credentials.databricks"
|
||||
]
|
||||
# WARNING+ would re-introduce the stderr traceback.
|
||||
warnings_or_louder = [r for r in module_records if r.levelno >= logging.WARNING]
|
||||
assert warnings_or_louder == [], (
|
||||
f"Expected no WARNING+ records, got "
|
||||
f"{[(r.levelname, r.getMessage()) for r in warnings_or_louder]}."
|
||||
)
|
||||
|
||||
# INFO record preserves the diagnostic for cli-*.log post-mortem.
|
||||
info_records = [
|
||||
r
|
||||
for r in module_records
|
||||
if r.levelno == logging.INFO and "Config(profile=" in r.getMessage()
|
||||
]
|
||||
assert len(info_records) == 1, (
|
||||
f"Expected one INFO record, got {len(info_records)}. "
|
||||
f"Records: {[(r.levelname, r.getMessage()) for r in module_records]}."
|
||||
)
|
||||
# exc_info is (type, value, traceback); [2] must be real, else the
|
||||
# frames won't render in the log.
|
||||
exc_info = info_records[0].exc_info
|
||||
assert exc_info is not None and exc_info[2] is not None, (
|
||||
"INFO record must carry exc_info with a real traceback."
|
||||
)
|
||||
|
||||
|
||||
def test_honors_databricks_config_file_env(
|
||||
monkeypatch: pytest.MonkeyPatch, tmp_path: Path
|
||||
) -> None:
|
||||
# Two cfg files in different locations; we point
|
||||
# DATABRICKS_CONFIG_FILE at the second one and expect that one to
|
||||
# be used. This proves the env var overrides the default
|
||||
# ~/.databrickscfg path.
|
||||
decoy_dir = tmp_path / "decoy"
|
||||
decoy_dir.mkdir()
|
||||
real_dir = tmp_path / "real"
|
||||
real_dir.mkdir()
|
||||
decoy = _write_cfg(
|
||||
decoy_dir,
|
||||
"[DEFAULT]\nhost = https://decoy.example.com\ntoken = decoy-token\n",
|
||||
)
|
||||
real = _write_cfg(
|
||||
real_dir,
|
||||
"[DEFAULT]\nhost = https://real.example.com\ntoken = real-token\n",
|
||||
)
|
||||
assert decoy.exists() # sanity check: both files were written
|
||||
monkeypatch.setenv("DATABRICKS_CONFIG_FILE", str(real))
|
||||
|
||||
creds = resolve_databricks_workspace(profile=None)
|
||||
|
||||
# The resolver read from `real`, not `decoy`. If this assertion
|
||||
# fails on the decoy host, the env var override is broken.
|
||||
assert creds.host == "https://real.example.com"
|
||||
assert creds.token == "real-token"
|
||||
Reference in New Issue
Block a user