scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-11 22:51:28 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 154

Fix Databricks SDK test failures: Mock databricks.sdk module correctly

Browse source 
- Fix test_sdk_partner_registered: Use patch.dict to mock databricks.sdk module in sys.modules
- Fix test_sdk_fallback_when_no_credentials: Use patch.dict to mock databricks.sdk module in sys.modules
- The code imports 'from databricks.sdk import ...' inside functions, and the SDK is not installed in test environment
- Mocking sys.modules allows the import to succeed and tests still verify the correct behavior:
  * useragent.with_partner('litellm') is called
  * SDK fallback authentication works correctly
This commit is contained in:
Alexsander Hamir 2026-01-01 10:47:54 -08:00
parent 22a2dacd9c
commit 5c7815c4c2
1 changed files with 40 additions and 18 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

58
tests/test_litellm/llms/databricks/test_databricks_partner_integration.py
View file

@ -349,15 +349,26 @@ class TestSDKPartnerTelemetry:
}
mock_useragent = MagicMock()
with patch("databricks.sdk.WorkspaceClient", return_value=mock_workspace_client):
with patch("databricks.sdk.useragent", mock_useragent):
databricks_base._get_databricks_credentials(
api_key=None,
api_base=None,
headers=None,
)
# 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,
)
mock_useragent.with_partner.assert_called_once_with("litellm")
# Verify that partner telemetry registration was called correctly
mock_useragent.with_partner.assert_called_once_with("litellm")
class TestUserAgentFromEnvironment:
@ -592,17 +603,28 @@ class TestAuthenticationPriority:
"Authorization": "Bearer sdk-token"
}
with patch("databricks.sdk.WorkspaceClient", return_value=mock_workspace_client):
with patch("databricks.sdk.useragent"):
api_base, headers = databricks_base.databricks_validate_environment(
api_key=None,
api_base=None,
endpoint_type="chat_completions",
custom_endpoint=False,
headers=None,
)
# 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,
)
assert "Authorization" in headers
# Verify that SDK authentication was used (headers contain Authorization)
assert "Authorization" in headers
class TestEndpointURLConstruction: