omnigent-ai--omnigent/tests/runtime/credentials/test_databricks.py

"""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"