Skip to content

Commit 01c3c98

Browse files
renaudhartert-dbrenaudhartert-db
committed
add tests
1 parent fe9877a commit 01c3c98

File tree

2 files changed

+153
-5
lines changed

2 files changed

+153
-5
lines changed

‎databricks/sdk/credentials_provider.py‎

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -199,8 +199,8 @@ def external_browser(cfg: 'Config') -> Optional[CredentialsProvider]:
199199
if not client_id:
200200
client_id = 'databricks-cli'
201201

202-
# Load cached credentials from disk if they exist.
203-
# Note that these are local to the Python SDK and not reused by other SDKs.
202+
# Load cached credentials from disk if they exist. Note that these are
203+
# local to the Python SDK and not reused by other SDKs.
204204
oidc_endpoints = cfg.oidc_endpoints
205205
redirect_url = 'http://localhost:8020'
206206
token_cache = TokenCache(host=cfg.host,
@@ -210,10 +210,11 @@ def external_browser(cfg: 'Config') -> Optional[CredentialsProvider]:
210210
redirect_url=redirect_url)
211211
credentials = token_cache.load()
212212
if credentials:
213-
# Force a refresh in case the loaded credentials are expired.
214-
# If the refresh fails, rather than throw exception we will initiate a new OAuth login flow.
215213
try:
216-
credentials.token()
214+
# Force a refresh in case the loaded credentials are expired.
215+
# If the refresh fails, rather than throw exception we will
216+
# initiate a new OAuth login flow.
217+
credentials.token() # force a token refresh
217218
return credentials(cfg)
218219
# TODO: we should ideally use more specific exceptions.
219220
except Exception as e:

‎tests/test_credentials_provider.py‎

Lines changed: 147 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,147 @@
1+
import pytest
2+
3+
from unittest.mock import Mock
4+
from databricks.sdk.credentials_provider import external_browser
5+
6+
def test_external_browser_refresh_success(mocker):
7+
"""Tests successful refresh of existing credentials."""
8+
9+
# 1. Mock Config
10+
mock_cfg = Mock()
11+
mock_cfg.auth_type = 'external-browser'
12+
mock_cfg.host = 'test-host'
13+
mock_cfg.oidc_endpoints = {'token_endpoint': 'test-token-endpoint'}
14+
mock_cfg.client_id = 'test-client-id' # Or use azure_client_id
15+
mock_cfg.client_secret = 'test-client-secret' # Or use azure_client_secret
16+
17+
# 2. Mock TokenCache
18+
mock_token_cache = Mock()
19+
mock_session_credentials = Mock()
20+
mock_session_credentials.token.return_value = "valid_token" # Simulate successful refresh
21+
mock_token_cache.load.return_value = mock_session_credentials
22+
23+
mock_credentials_provider = Mock()
24+
mock_session_credentials.return_value = mock_credentials_provider
25+
26+
# 3. Patch TokenCache (no need to mock OAuthClient in this case)
27+
mocker.patch('databricks.sdk.credentials_provider.TokenCache', return_value=mock_token_cache)
28+
29+
# 4. Call the function
30+
result = external_browser(mock_cfg)
31+
32+
# 5. Assertions
33+
mock_token_cache.load.assert_called_once()
34+
mock_session_credentials.token.assert_called_once() # Verify token refresh was attempted
35+
assert result == mock_credentials_provider
36+
37+
38+
def test_external_browser_refresh_failure_new_oauth_flow(mocker):
39+
"""Tests failed refresh, triggering a new OAuth flow."""
40+
41+
# 1. Mock Config
42+
mock_cfg = Mock()
43+
mock_cfg.auth_type = 'external-browser'
44+
mock_cfg.host = 'test-host'
45+
mock_cfg.oidc_endpoints = {'token_endpoint': 'test-token-endpoint'}
46+
mock_cfg.client_id = 'test-client-id'
47+
mock_cfg.client_secret = 'test-client-secret'
48+
49+
# 2. Mock TokenCache
50+
mock_token_cache = Mock()
51+
mock_session_credentials = Mock()
52+
mock_session_credentials.token.side_effect = Exception("Simulated refresh error") # Simulate a failed refresh
53+
mock_token_cache.load.return_value = mock_session_credentials
54+
55+
mock_credentials_provider = Mock()
56+
mock_session_credentials.return_value = mock_credentials_provider
57+
58+
# 3. Mock OAuthClient
59+
mock_oauth_client = Mock()
60+
mock_consent = Mock()
61+
mock_consent.launch_external_browser.return_value = mock_session_credentials # Simulate successful OAuth flow
62+
mock_oauth_client.initiate_consent.return_value = mock_consent
63+
64+
# 4. Patch TokenCache and OAuthClient
65+
mocker.patch('databricks.sdk.credentials_provider.TokenCache', return_value=mock_token_cache)
66+
mocker.patch('databricks.sdk.credentials_provider.OAuthClient', return_value=mock_oauth_client)
67+
68+
# 5. Call the function
69+
result = external_browser(mock_cfg)
70+
71+
# 6. Assertions
72+
mock_token_cache.load.assert_called_once()
73+
mock_session_credentials.token.assert_called_once() # Refresh attempt
74+
mock_oauth_client.initiate_consent.assert_called_once()
75+
mock_consent.launch_external_browser.assert_called_once()
76+
mock_token_cache.save.assert_called_once_with(mock_session_credentials)
77+
assert result == mock_credentials_provider
78+
79+
80+
def test_external_browser_no_cached_credentials(mocker):
81+
"""Tests the case where there are no cached credentials, initiating a new OAuth flow."""
82+
83+
# 1. Mock Config
84+
mock_cfg = Mock()
85+
mock_cfg.auth_type = 'external-browser'
86+
mock_cfg.host = 'test-host'
87+
mock_cfg.oidc_endpoints = {'token_endpoint': 'test-token-endpoint'}
88+
mock_cfg.client_id = 'test-client-id'
89+
mock_cfg.client_secret = 'test-client-secret'
90+
91+
# 2. Mock TokenCache
92+
mock_token_cache = Mock()
93+
mock_token_cache.load.return_value = None # No cached credentials
94+
95+
mock_session_credentials = lambda c: "new_credentials"
96+
97+
# 3. Mock OAuthClient
98+
mock_consent = Mock()
99+
mock_consent.launch_external_browser.return_value = mock_session_credentials
100+
mock_oauth_client = Mock()
101+
mock_oauth_client.initiate_consent.return_value = mock_consent
102+
103+
# 4. Patch TokenCache and OAuthClient
104+
mocker.patch('databricks.sdk.credentials_provider.TokenCache', return_value=mock_token_cache)
105+
mocker.patch('databricks.sdk.credentials_provider.OAuthClient', return_value=mock_oauth_client)
106+
107+
# 5. Call the function
108+
result = external_browser(mock_cfg)
109+
110+
# 6. Assertions
111+
mock_token_cache.load.assert_called_once()
112+
mock_oauth_client.initiate_consent.assert_called_once()
113+
mock_consent.launch_external_browser.assert_called_once()
114+
mock_token_cache.save.assert_called_once_with(mock_session_credentials)
115+
assert result == "new_credentials"
116+
117+
118+
def test_external_browser_consent_fails(mocker):
119+
"""Tests the case where OAuth consent initiation fails."""
120+
121+
# 1. Mock Config
122+
mock_cfg = Mock()
123+
mock_cfg.auth_type = 'external-browser'
124+
mock_cfg.host = 'test-host'
125+
mock_cfg.oidc_endpoints = {'token_endpoint': 'test-token-endpoint'}
126+
mock_cfg.client_id = 'test-client-id'
127+
mock_cfg.client_secret = 'test-client-secret'
128+
129+
# 2. Mock TokenCache
130+
mock_token_cache = Mock()
131+
mock_token_cache.load.return_value = None # No cached credentials
132+
133+
# 3. Mock OAuthClient
134+
mock_oauth_client = Mock()
135+
mock_oauth_client.initiate_consent.return_value = None # Simulate consent failure
136+
137+
# 4. Patch TokenCache and OAuthClient
138+
mocker.patch('databricks.sdk.credentials_provider.TokenCache', return_value=mock_token_cache)
139+
mocker.patch('databricks.sdk.credentials_provider.OAuthClient', return_value=mock_oauth_client)
140+
141+
# 5. Call the function
142+
result = external_browser(mock_cfg)
143+
144+
# 6. Assertions
145+
mock_token_cache.load.assert_called_once()
146+
mock_oauth_client.initiate_consent.assert_called_once()
147+
assert result is None

0 commit comments

Comments
 (0)