From 5c7815c4c2d8fb36d70cc2849128ed41b9b4f143 Mon Sep 17 00:00:00 2001
From: Alexsander Hamir <alexsanderhamirgomesbaptista@gmail.com>
Date: Thu, 1 Jan 2026 10:47:54 -0800
Subject: [PATCH] Fix Databricks SDK test failures: Mock databricks.sdk module
 correctly

- 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
---
 .../test_databricks_partner_integration.py    | 58 +++++++++++++------
 1 file changed, 40 insertions(+), 18 deletions(-)

diff --git a/tests/test_litellm/llms/databricks/test_databricks_partner_integration.py b/tests/test_litellm/llms/databricks/test_databricks_partner_integration.py
index 9bcd62c14b0..800066ac5bf 100644
--- a/tests/test_litellm/llms/databricks/test_databricks_partner_integration.py
+++ b/tests/test_litellm/llms/databricks/test_databricks_partner_integration.py
@@ -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: