Fix tests

Author: aravind-segu

tests/test_notebook_oauth.py

@@ -1,14 +1,18 @@
 """Tests for runtime OAuth authentication in notebook environments."""
 
 import os
+import sys
+import types
 from datetime import datetime, timedelta
-from unittest.mock import patch
+from typing import Dict
 
 import pytest
 
 from databricks.sdk import oauth
 from databricks.sdk.config import Config
-from databricks.sdk.credentials_provider import (DefaultCredentials,
+from databricks.sdk.credentials_provider import (CredentialsProvider,
+                                                 CredentialsStrategy,
+                                                 DefaultCredentials,
                                                  runtime_oauth)
 
 
@@ -24,14 +28,26 @@ def mock_runtime_env(monkeypatch):
 @pytest.fixture
 def mock_runtime_native_auth():
     """Mock the runtime_native_auth to return a valid credentials provider."""
-    with patch("databricks.sdk.runtime.init_runtime_native_auth") as mock_auth:
+    fake_runtime = types.ModuleType("databricks.sdk.runtime")
 
+    def fake_init_runtime_native_auth():
         def inner():
             return {"Authorization": "Bearer test-notebook-pat-token"}
 
-        mock_auth.return_value = ("https://test.cloud.databricks.com", inner)
-        mock_auth.__name__ = "init_runtime_native_auth"
-        yield mock_auth
+        return "https://test.cloud.databricks.com", inner
+
+    def fake_init_runtime_legacy_auth():
+        pass
+
+    def fake_init_runtime_repl_auth():
+        pass
+
+    fake_runtime.init_runtime_native_auth = fake_init_runtime_native_auth
+    fake_runtime.init_runtime_legacy_auth = fake_init_runtime_legacy_auth
+    fake_runtime.init_runtime_repl_auth = fake_init_runtime_repl_auth
+
+    sys.modules["databricks.sdk.runtime"] = fake_runtime
+    yield
 
 
 @pytest.fixture
@@ -48,6 +64,17 @@ def mock_pat_exchange(mocker):
     return mock_exchange
 
 
+class MockCredentialsStrategy(CredentialsStrategy):
+    def auth_type(self) -> str:
+        return "mock_credentials_strategy"
+
+    def __call__(self, cfg) -> CredentialsProvider:
+        def credentials_provider() -> Dict[str, str]:
+            return {"Authorization": "Bearer: no_token"}
+
+        return credentials_provider
+
+
 @pytest.mark.parametrize(
     "scopes,auth_details",
     [
@@ -61,8 +88,12 @@ def test_runtime_oauth_success_scenarios(
     mock_runtime_env, mock_runtime_native_auth, mock_pat_exchange, scopes, auth_details
 ):
     """Test runtime-oauth works correctly in various valid configurations."""
-    cfg = Config(host="https://test.cloud.databricks.com", scopes=scopes, authorization_details=auth_details)
-
+    cfg = Config(
+        host="https://test.cloud.databricks.com",
+        scopes=scopes,
+        authorization_details=auth_details,
+        credentials_strategy=MockCredentialsStrategy(),
+    )
     creds_provider = runtime_oauth(cfg)
 
     assert creds_provider is not None