scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-27 01:22:18 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 80
litellm/tests/unit/llms/databricks/test_databricks_partner_integration.py
yuneng-jiang 5e6dc89ba1
test: move tests/test_litellm/llms into tests/unit/llms (#43191) 
* ci: run the unit_selection.sh shard files on every event instead of only fork pull requests

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* ci: rename fork-flag to unit-flag now that it applies on every event

* test: move tests/test_litellm root and small trees into tests/unit

Pure renames, no content changes. Follow-up commits in this PR fix
references, merge the three files that already existed in tests/unit,
keep live-provider tests in tests/test_litellm and wire CI.

* test: carry tests/test_litellm conftest isolation into tests/unit

Callback lists, routing fallbacks, cached HTTP clients, logger state, AWS,
proxy-URL and keychain env, and session-end client cleanup now reset for
unit tests too. The environment isolation owns its MonkeyPatch so a test's
own monkeypatch is undone before the model-cost teardown runs.

* test: merge, split and prune the moved root and small-tree tests

Merge batches/test_batch_utils.py and the chat_completions and messages
dispatch tests into the files that already existed in tests/unit. Keep
the live Gemini interactions tests, the async image-fetch format test and
the OpenAI embedding scorer test in tests/test_litellm since they need
real network or keys. Put test_router.py under tests/unit/test_router so
the existing package no longer shadows it. Delete eight tests the audit
found superseded by stronger ones kept in this move.

* ci: run the moved root and small-tree tests under their legacy flags

Add the misc and responses-caching-types flags to unit_selection.sh and
CircleCI, extend enterprise-routing and mcp-integration, and point the
legacy GHA shards, Makefile, redis-compat workflow, merge smoke manifest
and change classifier at the new paths.

* test: make the new tests/unit directories packages

tests/unit/test_package_layout.py requires every directory to carry an
__init__.py, and without one the moved and retained
test_litellm_responses_bridge.py modules collide on import.

* test: scope the unit socket block to tests/unit in shared sessions

The GHA shards collect the legacy test-path and the unit selection in one
pytest session. The unit conftest's loopback-only block leaked into legacy
modules that reach the network at import. The legacy conftest now lifts the
restriction at collect and setup time, and the unit conftest re-applies it
when collecting its own modules.

* test: move tests/test_litellm/llms into tests/unit/llms

Rename-only. Moves the provider tests and the fine-tuning fixtures they
load, mirroring the old paths. Follow-up commits merge, split and wire them.

* test: merge, split and prune the moved llms tests

Merges the Databricks chat transformation tests into the existing unit
file, keeps the tests that need real keys or the network in
tests/test_litellm, deletes the audited tests a stronger unit test
already covers, and points imports at tests.unit.llms.

* ci: run the moved llms tests under their legacy flags

The Vertex AI and All Other Providers shards keep their legacy test-path
for the retained files and add the llm-vertex-ai and llm-other-providers
unit selections. CircleCI gets matching unit jobs.

* test: make the tests/unit/llms directories packages

Adds __init__.py to the moved dirs and drops the legacy ones whose
directories no longer hold tests.

* test: drop script runners and path hacks the llms split left dangling

The __main__ runners in the split openai_like files and the Databricks e2e
runner called tests that now live in the other half of the split or were
deleted. The retained legacy halves also no longer need sys.path edits.

* test: give the shard-script tests their own GITHUB_OUTPUT

They only passed where the runner set it. The CircleCI unit job's env
allowlist drops it, so the script's redirect failed there.

* test: point the router and module-deletion checks at tests/unit

router_code_coverage and code_qa_check_tests only searched tests/test_litellm,
so the moved router tests no longer counted. The two silent-experiment tests
the audit deleted were the only direct callers of those methods; they are
replaced with tests that assert the forwarded shadow request and the
recursion guard.

* test: keep the Databricks manual e2e runner and fix the SageMaker Nova run path

The Databricks e2e file is a manual script whose main() calls the tests
that were pruned, so pruning them broke the documented run. It is back to
its main version. The SageMaker Nova docstring now points at the file's
real location in tests/local_testing.

* test: keep the job's UNIT_FLAG out of the shard-script tests

---------

Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
2026-09-25 12:43:23 -07:00

763 lines
30 KiB
Python
Raw Blame History

"""
Unit Tests for Databricks Partner Integration Features
=======================================================
These tests are designed for automated CI/CD pipelines and do NOT require
real Databricks credentials. All external calls are mocked.
For integration tests that use real Databricks credentials, see:
test_databricks_integration.py
Features Tested:
- User-Agent building with partner prefixing (Databricks partner telemetry)
- Token/sensitive data redaction for secure logging
- OAuth M2M (Machine-to-Machine) authentication flow
- Databricks SDK partner telemetry registration
- Authentication priority (OAuth M2M > PAT > SDK)
Run with:
pytest test_databricks_partner_integration.py -v
These tests align with Databricks Partner Architecture best practices:
https://github.com/databrickslabs/partner-architecture
"""
import json
import sys
import pytest
from unittest.mock import MagicMock, patch, Mock
from litellm.llms.databricks.common_utils import DatabricksBase, DatabricksException
class TestBuildUserAgent:
"""Test cases for User-Agent string building."""
def test_default_user_agent(self):
"""No custom user agent returns litellm/{version}."""
ua = DatabricksBase._build_user_agent(None)
assert ua.startswith("litellm/")
assert "_" not in ua.split("/")[0]
def test_custom_user_agent_with_version(self):
"""Custom user agent with version extracts partner name."""
ua = DatabricksBase._build_user_agent("mycompany/1.0.0")
assert ua.startswith("mycompany_litellm/")
# Verify the version is litellm's, not the custom one
assert "/1.0.0" not in ua or "mycompany_litellm/1.0.0" not in ua
def test_custom_user_agent_without_version(self):
"""Custom user agent without version still works."""
ua = DatabricksBase._build_user_agent("mycompany")
assert ua.startswith("mycompany_litellm/")
def test_custom_user_agent_with_underscore(self):
"""Partner names with underscores are preserved."""
ua = DatabricksBase._build_user_agent("my_company/2.0.0")
assert ua.startswith("my_company_litellm/")
def test_custom_user_agent_with_hyphen(self):
"""Partner names with hyphens are preserved."""
ua = DatabricksBase._build_user_agent("my-company/2.0.0")
assert ua.startswith("my-company_litellm/")
def test_custom_user_agent_ignores_custom_version(self):
"""Custom version is ignored, litellm version is used."""
ua = DatabricksBase._build_user_agent("partner/99.99.99")
parts = ua.split("/")
assert parts[0] == "partner_litellm"
assert parts[1] != "99.99.99"
def test_empty_string_returns_default(self):
"""Empty string returns default user agent."""
ua = DatabricksBase._build_user_agent("")
assert ua.startswith("litellm/")
assert "_" not in ua.split("/")[0]
def test_whitespace_only_returns_default(self):
"""Whitespace-only string returns default user agent."""
ua = DatabricksBase._build_user_agent(" ")
assert ua.startswith("litellm/")
assert "_" not in ua.split("/")[0]
def test_invalid_partner_name_returns_default(self):
"""Invalid partner names (special chars) return default."""
ua = DatabricksBase._build_user_agent("my@company/1.0.0")
assert ua.startswith("litellm/")
def test_partner_with_numbers(self):
"""Partner names with numbers work."""
ua = DatabricksBase._build_user_agent("company123/1.0.0")
assert ua.startswith("company123_litellm/")
class TestRedactSensitiveData:
"""Test cases for sensitive data redaction."""
def test_redact_bearer_token_in_string(self):
"""Bearer tokens are redacted in strings."""
result = DatabricksBase.redact_sensitive_data("Bearer dapi12345abcdef")
assert "dapi12345abcdef" not in result
assert "[REDACTED]" in result
def test_redact_dict_with_authorization(self):
"""Dict with authorization key is redacted."""
data = {"Authorization": "Bearer secret123", "other": "value"}
result = DatabricksBase.redact_sensitive_data(data)
assert result["Authorization"] == "[REDACTED]"
assert result["other"] == "value"
def test_redact_nested_dict(self):
"""Nested dicts with sensitive keys are redacted."""
data = {"config": {"api_key": "secret", "name": "test"}}
result = DatabricksBase.redact_sensitive_data(data)
assert result["config"]["api_key"] == "[REDACTED]"
assert result["config"]["name"] == "test"
def test_redact_pat_token(self):
"""Databricks PAT tokens are redacted."""
test_token = "dapiTESTTOKENFAKEVALUEFORTESTINGPURPOSESONLY123"
result = DatabricksBase.redact_sensitive_data(f"Using token {test_token}")
assert test_token not in result
assert "[REDACTED_PAT]" in result
def test_redact_client_secret(self):
"""Client secrets are redacted."""
data = {"client_secret": "my-super-secret-value"}
result = DatabricksBase.redact_sensitive_data(data)
assert result["client_secret"] == "[REDACTED]"
def test_redact_list_of_dicts(self):
"""Lists containing dicts with sensitive data are redacted."""
data = [{"api_key": "secret1"}, {"name": "test"}]
result = DatabricksBase.redact_sensitive_data(data)
assert result[0]["api_key"] == "[REDACTED]"
assert result[1]["name"] == "test"
def test_redact_none_returns_none(self):
"""None input returns None."""
assert DatabricksBase.redact_sensitive_data(None) is None
def test_redact_preserves_non_sensitive_data(self):
"""Non-sensitive data is preserved."""
data = {"model": "dbrx", "temperature": 0.7, "messages": ["hello"]}
result = DatabricksBase.redact_sensitive_data(data)
assert result == data
class TestRedactHeadersForLogging:
"""Test cases for header redaction."""
def test_authorization_header_partially_shown(self):
"""Authorization header shows first 8 chars then redacts."""
headers = {"Authorization": "Bearer dapi123456789abcdef"}
result = DatabricksBase.redact_headers_for_logging(headers)
assert result["Authorization"].startswith("Bearer d")
assert "[REDACTED]" in result["Authorization"]
def test_short_authorization_header_fully_redacted(self):
"""Short authorization values are fully redacted."""
headers = {"Authorization": "short"}
result = DatabricksBase.redact_headers_for_logging(headers)
assert result["Authorization"] == "[REDACTED]"
def test_non_sensitive_headers_preserved(self):
"""Non-sensitive headers are not modified."""
headers = {"Content-Type": "application/json", "User-Agent": "test/1.0"}
result = DatabricksBase.redact_headers_for_logging(headers)
assert result["Content-Type"] == "application/json"
assert result["User-Agent"] == "test/1.0"
def test_empty_headers_returns_empty(self):
"""Empty headers dict returns empty dict."""
assert DatabricksBase.redact_headers_for_logging({}) == {}
def test_none_headers_returns_empty(self):
"""None headers returns empty dict."""
assert DatabricksBase.redact_headers_for_logging(None) == {}
def test_x_api_key_header_redacted(self):
"""X-API-Key header is redacted."""
headers = {"X-API-Key": "my-api-key-12345"}
result = DatabricksBase.redact_headers_for_logging(headers)
assert "[REDACTED]" in result["X-API-Key"]
class TestOAuthM2M:
"""Test cases for OAuth M2M authentication."""
def test_oauth_m2m_token_success(self):
"""OAuth M2M token is successfully obtained."""
databricks_base = DatabricksBase()
mock_response = Mock()
mock_response.status_code = 200
mock_response.json.return_value = {"access_token": "test-access-token"}
with patch("requests.post", return_value=mock_response) as mock_post:
token = databricks_base._get_oauth_m2m_token(
api_base="https://adb-123.azuredatabricks.net/serving-endpoints",
client_id="test-client-id",
client_secret="test-client-secret",
)
assert token == "test-access-token"
mock_post.assert_called_once()
call_args = mock_post.call_args
assert "oidc/v1/token" in call_args[0][0]
assert call_args[1]["data"]["grant_type"] == "client_credentials"
def test_oauth_m2m_token_failure(self):
"""OAuth M2M raises exception on failure."""
databricks_base = DatabricksBase()
mock_response = Mock()
mock_response.status_code = 401
mock_response.text = "Unauthorized"
with patch("requests.post", return_value=mock_response):
with pytest.raises(DatabricksException) as exc_info:
databricks_base._get_oauth_m2m_token(
api_base="https://adb-123.azuredatabricks.net",
client_id="bad-client-id",
client_secret="bad-secret",
)
assert exc_info.value.status_code == 401
def test_oauth_m2m_strips_serving_endpoints(self):
"""OAuth M2M correctly strips /serving-endpoints from URL."""
databricks_base = DatabricksBase()
mock_response = Mock()
mock_response.status_code = 200
mock_response.json.return_value = {"access_token": "token"}
with patch("requests.post", return_value=mock_response) as mock_post:
databricks_base._get_oauth_m2m_token(
api_base="https://adb-123.azuredatabricks.net/serving-endpoints",
client_id="id",
client_secret="secret",
)
call_url = mock_post.call_args[0][0]
assert "/serving-endpoints" not in call_url
assert call_url == "https://adb-123.azuredatabricks.net/oidc/v1/token"
def test_oauth_m2m_strips_ai_gateway_path(self):
"""OAuth M2M derives the token URL from the workspace origin."""
databricks_base = DatabricksBase()
mock_response = Mock()
mock_response.status_code = 200
mock_response.json.return_value = {"access_token": "token"}
with patch("requests.post", return_value=mock_response) as mock_post:
databricks_base._get_oauth_m2m_token(
api_base="https://adb-123.azuredatabricks.net/ai-gateway/mlflow/v1",
client_id="id",
client_secret="secret",
)
call_url = mock_post.call_args[0][0]
assert call_url == "https://adb-123.azuredatabricks.net/oidc/v1/token"
class TestValidateEnvironmentWithOAuth:
"""Test OAuth M2M is used when credentials are available."""
def test_oauth_used_when_credentials_set(self, monkeypatch):
"""OAuth M2M is used when client_id and client_secret are set."""
monkeypatch.setenv("DATABRICKS_CLIENT_ID", "test-client-id")
monkeypatch.setenv("DATABRICKS_CLIENT_SECRET", "test-secret")
monkeypatch.setenv(
"DATABRICKS_API_BASE", "https://adb-123.net/serving-endpoints"
)
databricks_base = DatabricksBase()
with patch.object(
databricks_base, "_get_oauth_m2m_token", return_value="oauth-token"
) as mock_oauth:
api_base, headers = databricks_base.databricks_validate_environment(
api_key=None,
api_base=None,
endpoint_type="chat_completions",
custom_endpoint=False,
headers=None,
)
mock_oauth.assert_called_once()
assert headers["Authorization"] == "Bearer oauth-token"
def test_pat_used_when_api_key_set(self, monkeypatch):
"""PAT is used when api_key is provided."""
monkeypatch.delenv("DATABRICKS_CLIENT_ID", raising=False)
monkeypatch.delenv("DATABRICKS_CLIENT_SECRET", raising=False)
databricks_base = DatabricksBase()
api_base, headers = databricks_base.databricks_validate_environment(
api_key="dapi-test-key",
api_base="https://adb-123.net/serving-endpoints",
endpoint_type="chat_completions",
custom_endpoint=False,
headers=None,
)
assert headers["Authorization"] == "Bearer dapi-test-key"
class TestValidateEnvironmentUserAgent:
"""Test User-Agent is correctly set in validate_environment."""
def test_default_user_agent(self, monkeypatch):
"""Default user agent is set when no custom agent provided."""
monkeypatch.delenv("DATABRICKS_USER_AGENT", raising=False)
monkeypatch.delenv("LITELLM_USER_AGENT", raising=False)
databricks_base = DatabricksBase()
api_base, headers = databricks_base.databricks_validate_environment(
api_key="test-key",
api_base="https://adb-123.net/serving-endpoints",
endpoint_type="chat_completions",
custom_endpoint=False,
headers=None,
custom_user_agent=None,
)
assert headers["User-Agent"].startswith("litellm/")
assert "_" not in headers["User-Agent"].split("/")[0]
def test_custom_user_agent_via_param(self, monkeypatch):
"""Custom user agent is prefixed when passed as parameter."""
databricks_base = DatabricksBase()
api_base, headers = databricks_base.databricks_validate_environment(
api_key="test-key",
api_base="https://adb-123.net/serving-endpoints",
endpoint_type="chat_completions",
custom_endpoint=False,
headers=None,
custom_user_agent="mycompany/1.0.0",
)
assert headers["User-Agent"].startswith("mycompany_litellm/")
class TestSDKPartnerTelemetry:
"""Test that SDK partner telemetry is registered."""
def test_sdk_partner_registered(self):
"""useragent.with_partner is called when using SDK."""
databricks_base = DatabricksBase()
mock_workspace_client = MagicMock()
mock_workspace_client.config.host = "https://adb-123.net"
mock_workspace_client.config.authenticate.return_value = {
"Authorization": "Bearer token"
}
mock_useragent = MagicMock()
# Create a mock databricks.sdk module to simulate the SDK being available
# This allows us to test the partner telemetry registration without requiring
# the actual databricks-sdk package to be installed
mock_sdk_module = MagicMock()
mock_sdk_module.WorkspaceClient = MagicMock(return_value=mock_workspace_client)
mock_sdk_module.useragent = mock_useragent
# Mock both databricks and databricks.sdk modules to ensure the import works
with patch.dict(
sys.modules, {"databricks": MagicMock(), "databricks.sdk": mock_sdk_module}
):
databricks_base._get_databricks_credentials(
api_key=None,
api_base=None,
headers=None,
)
# Verify that partner telemetry registration was called correctly
mock_useragent.with_partner.assert_called_once_with("litellm")
class TestUserAgentFromEnvironment:
"""Test User-Agent is correctly picked up from environment variables."""
def test_user_agent_from_databricks_env_var(self, monkeypatch):
"""DATABRICKS_USER_AGENT environment variable is used."""
monkeypatch.setenv("DATABRICKS_USER_AGENT", "envpartner")
monkeypatch.delenv("LITELLM_USER_AGENT", raising=False)
databricks_base = DatabricksBase()
api_base, headers = databricks_base.databricks_validate_environment(
api_key="test-key",
api_base="https://adb-123.net/serving-endpoints",
endpoint_type="chat_completions",
custom_endpoint=False,
headers=None,
custom_user_agent="envpartner", # Simulating what transformation.py passes
)
assert headers["User-Agent"].startswith("envpartner_litellm/")
def test_custom_param_takes_precedence(self, monkeypatch):
"""Custom user_agent parameter takes precedence over environment."""
monkeypatch.setenv("DATABRICKS_USER_AGENT", "envpartner")
databricks_base = DatabricksBase()
api_base, headers = databricks_base.databricks_validate_environment(
api_key="test-key",
api_base="https://adb-123.net/serving-endpoints",
endpoint_type="chat_completions",
custom_endpoint=False,
headers=None,
custom_user_agent="parampartner/1.0.0",
)
assert headers["User-Agent"].startswith("parampartner_litellm/")
class TestLiteLLMCompletionUserAgent:
"""Test User-Agent is correctly passed through LiteLLM completion calls."""
def test_completion_passes_user_agent_to_headers(self):
"""litellm.completion() correctly passes user_agent to request headers."""
from litellm.llms.databricks.chat.transformation import DatabricksConfig
config = DatabricksConfig()
optional_params = {"user_agent": "testpartner/1.0.0"}
# Mock the validation to capture what headers are set
with patch.object(
config,
"databricks_validate_environment",
return_value=(
"https://test.net/serving-endpoints/chat/completions",
{
"Authorization": "Bearer test",
"User-Agent": "testpartner_litellm/1.0.0",
},
),
) as mock_validate:
result = config.validate_environment(
headers={},
model="databricks/test-model",
messages=[],
optional_params=optional_params,
litellm_params={},
api_key="test-key",
api_base="https://test.net/serving-endpoints",
)
# Verify user_agent was passed to databricks_validate_environment
mock_validate.assert_called_once()
call_kwargs = mock_validate.call_args[1]
assert call_kwargs.get("custom_user_agent") == "testpartner/1.0.0"
def test_user_agent_removed_from_optional_params(self):
"""user_agent is removed from optional_params so it's not sent to API."""
from litellm.llms.databricks.chat.transformation import DatabricksConfig
config = DatabricksConfig()
optional_params = {
"user_agent": "testpartner/1.0.0",
"temperature": 0.7,
}
with patch.object(
config,
"databricks_validate_environment",
return_value=(
"https://test.net/chat/completions",
{"Authorization": "Bearer test", "User-Agent": "test"},
),
):
config.validate_environment(
headers={},
model="databricks/test-model",
messages=[],
optional_params=optional_params,
litellm_params={},
api_key="test-key",
api_base="https://test.net/serving-endpoints",
)
# user_agent should be removed from optional_params
assert "user_agent" not in optional_params
# Other params should remain
assert optional_params.get("temperature") == 0.7
class TestLiteLLMEmbeddingUserAgent:
"""Test User-Agent is correctly passed through LiteLLM embedding calls."""
def test_embedding_passes_user_agent_to_headers(self):
"""litellm.embedding() correctly passes user_agent to request headers."""
from litellm.llms.databricks.embed.handler import DatabricksEmbeddingHandler
handler = DatabricksEmbeddingHandler()
optional_params = {"user_agent": "embedpartner/1.0.0"}
with patch.object(
handler,
"databricks_validate_environment",
return_value=(
"https://test.net/serving-endpoints/embeddings",
{
"Authorization": "Bearer test",
"User-Agent": "embedpartner_litellm/1.0.0",
},
),
) as mock_validate:
with patch(
"litellm.llms.openai_like.embedding.handler.OpenAILikeEmbeddingHandler.embedding"
):
try:
handler.embedding(
model="databricks/test-model",
input=["test"],
timeout=30,
api_key="test-key",
api_base="https://test.net/serving-endpoints",
optional_params=optional_params,
)
except Exception:
pass # We just want to verify the mock was called
# Verify user_agent was passed
if mock_validate.called:
call_kwargs = mock_validate.call_args[1]
assert call_kwargs.get("custom_user_agent") == "embedpartner/1.0.0"
class TestAuthenticationPriority:
"""Test that authentication methods are used in correct priority order."""
def test_oauth_used_when_no_api_key_provided(self, monkeypatch):
"""OAuth M2M is used when OAuth creds are set and no api_key is provided."""
monkeypatch.setenv("DATABRICKS_CLIENT_ID", "oauth-client-id")
monkeypatch.setenv("DATABRICKS_CLIENT_SECRET", "oauth-secret")
monkeypatch.setenv("DATABRICKS_API_BASE", "https://test.net/serving-endpoints")
databricks_base = DatabricksBase()
with patch.object(
databricks_base, "_get_oauth_m2m_token", return_value="oauth-token"
) as mock_oauth:
api_base, headers = databricks_base.databricks_validate_environment(
api_key=None, # No PAT provided - OAuth should be used
api_base=None,
endpoint_type="chat_completions",
custom_endpoint=False,
headers=None,
)
# OAuth should be used
mock_oauth.assert_called_once()
assert headers["Authorization"] == "Bearer oauth-token"
def test_explicit_pat_takes_priority_over_oauth_env(self, monkeypatch):
"""Explicit api_key takes priority over OAuth token in final headers."""
monkeypatch.setenv("DATABRICKS_CLIENT_ID", "oauth-client-id")
monkeypatch.setenv("DATABRICKS_CLIENT_SECRET", "oauth-secret")
monkeypatch.setenv("DATABRICKS_API_BASE", "https://test.net/serving-endpoints")
databricks_base = DatabricksBase()
# Mock the OAuth call - it will be attempted but PAT should override
with patch.object(
databricks_base, "_get_oauth_m2m_token", return_value="oauth-token"
):
api_base, headers = databricks_base.databricks_validate_environment(
api_key="dapi-explicit-pat",
api_base=None,
endpoint_type="chat_completions",
custom_endpoint=False,
headers=None,
)
# PAT should override OAuth token since api_key was explicitly provided
assert headers["Authorization"] == "Bearer dapi-explicit-pat"
def test_pat_used_when_no_oauth_credentials(self, monkeypatch):
"""PAT is used when OAuth credentials are not set."""
monkeypatch.delenv("DATABRICKS_CLIENT_ID", raising=False)
monkeypatch.delenv("DATABRICKS_CLIENT_SECRET", raising=False)
databricks_base = DatabricksBase()
api_base, headers = databricks_base.databricks_validate_environment(
api_key="dapi-pat-token",
api_base="https://test.net/serving-endpoints",
endpoint_type="chat_completions",
custom_endpoint=False,
headers=None,
)
assert headers["Authorization"] == "Bearer dapi-pat-token"
def test_sdk_fallback_when_no_credentials(self, monkeypatch):
"""Databricks SDK is used when no API key or OAuth credentials."""
monkeypatch.delenv("DATABRICKS_CLIENT_ID", raising=False)
monkeypatch.delenv("DATABRICKS_CLIENT_SECRET", raising=False)
monkeypatch.delenv("DATABRICKS_API_KEY", raising=False)
databricks_base = DatabricksBase()
mock_workspace_client = MagicMock()
mock_workspace_client.config.host = "https://adb-123.net"
mock_workspace_client.config.authenticate.return_value = {
"Authorization": "Bearer sdk-token"
}
# Create a mock databricks.sdk module to simulate the SDK being available
# This allows us to test the SDK fallback authentication without requiring
# the actual databricks-sdk package to be installed
mock_sdk_module = MagicMock()
mock_sdk_module.WorkspaceClient = MagicMock(return_value=mock_workspace_client)
mock_sdk_module.useragent = MagicMock()
# Mock both databricks and databricks.sdk modules to ensure the import works
with patch.dict(
sys.modules, {"databricks": MagicMock(), "databricks.sdk": mock_sdk_module}
):
api_base, headers = databricks_base.databricks_validate_environment(
api_key=None,
api_base=None,
endpoint_type="chat_completions",
custom_endpoint=False,
headers=None,
)
# Verify that SDK authentication was used (headers contain Authorization)
assert "Authorization" in headers
class TestEndpointURLConstruction:
"""Test that endpoint URLs are correctly constructed."""
def test_chat_completions_endpoint(self, monkeypatch):
"""Chat completions endpoint is correctly appended."""
monkeypatch.delenv("DATABRICKS_CLIENT_ID", raising=False)
monkeypatch.delenv("DATABRICKS_CLIENT_SECRET", raising=False)
databricks_base = DatabricksBase()
api_base, headers = databricks_base.databricks_validate_environment(
api_key="test-key",
api_base="https://test.net/serving-endpoints",
endpoint_type="chat_completions",
custom_endpoint=False,
headers=None,
)
assert api_base.endswith("/chat/completions")
def test_chat_gateway_endpoint_for_unity_model_on_legacy_base(self, monkeypatch):
from litellm.llms.databricks.chat.transformation import DatabricksConfig
monkeypatch.delenv("DATABRICKS_CLIENT_ID", raising=False)
monkeypatch.delenv("DATABRICKS_CLIENT_SECRET", raising=False)
url = DatabricksConfig().get_complete_url(
api_base="https://test.net/serving-endpoints",
api_key="test-key",
model="system.ai.kimi-k3",
optional_params={},
litellm_params={},
)
assert url == "https://test.net/ai-gateway/mlflow/v1/chat/completions"
def test_chat_gateway_endpoint_preserves_explicit_gateway_base(self, monkeypatch):
from litellm.llms.databricks.chat.transformation import DatabricksConfig
monkeypatch.delenv("DATABRICKS_CLIENT_ID", raising=False)
monkeypatch.delenv("DATABRICKS_CLIENT_SECRET", raising=False)
url = DatabricksConfig().get_complete_url(
api_base="https://test.net/ai-gateway/mlflow/v1/",
api_key="test-key",
model="system.ai.kimi-k3",
optional_params={},
litellm_params={},
)
assert url == "https://test.net/ai-gateway/mlflow/v1/chat/completions"
def test_chat_gateway_preserves_unity_model_service_name_with_explicit_base(self, monkeypatch):
from litellm.llms.databricks.chat.transformation import DatabricksConfig
monkeypatch.delenv("DATABRICKS_CLIENT_ID", raising=False)
monkeypatch.delenv("DATABRICKS_CLIENT_SECRET", raising=False)
config = DatabricksConfig()
request = config.transform_request(
model="catalog.schema.kimi-k3",
messages=[{"role": "user", "content": "hello"}],
optional_params={},
litellm_params={},
headers={},
)
assert config.get_complete_url(
api_base="https://test.net/ai-gateway/mlflow/v1",
api_key="test-key",
model="catalog.schema.kimi-k3",
optional_params={},
litellm_params={},
) == "https://test.net/ai-gateway/mlflow/v1/chat/completions"
assert request["model"] == "catalog.schema.kimi-k3"
def test_chat_legacy_endpoint_remains_default(self, monkeypatch):
from litellm.llms.databricks.chat.transformation import DatabricksConfig
monkeypatch.delenv("DATABRICKS_CLIENT_ID", raising=False)
monkeypatch.delenv("DATABRICKS_CLIENT_SECRET", raising=False)
url = DatabricksConfig().get_complete_url(
api_base="https://test.net/serving-endpoints",
api_key="test-key",
model="databricks-kimi-k3",
optional_params={},
litellm_params={},
)
assert url == "https://test.net/serving-endpoints/chat/completions"
def test_embeddings_endpoint(self, monkeypatch):
"""Embeddings endpoint is correctly appended."""
monkeypatch.delenv("DATABRICKS_CLIENT_ID", raising=False)
monkeypatch.delenv("DATABRICKS_CLIENT_SECRET", raising=False)
databricks_base = DatabricksBase()
api_base, headers = databricks_base.databricks_validate_environment(
api_key="test-key",
api_base="https://test.net/serving-endpoints",
endpoint_type="embeddings",
custom_endpoint=False,
headers=None,
)
assert api_base.endswith("/embeddings")
def test_custom_endpoint_not_modified(self, monkeypatch):
"""Custom endpoints are not modified."""
monkeypatch.delenv("DATABRICKS_CLIENT_ID", raising=False)
monkeypatch.delenv("DATABRICKS_CLIENT_SECRET", raising=False)
databricks_base = DatabricksBase()
api_base, headers = databricks_base.databricks_validate_environment(
api_key="test-key",
api_base="https://test.net/custom/endpoint",
endpoint_type="chat_completions",
custom_endpoint=True,
headers=None,
)
assert api_base == "https://test.net/custom/endpoint"
Reference in a new issue View git blame Copy permalink