Skip to content

Commit 11fdf46

Browse files
[Internal] Add unit tests for external-browser authentication (#863)
## What changes are proposed in this pull request? This PR is a noop that adds unit tests for the "external-browser" credential strategy. ## How is this tested? This PR adds the unit test. Though, these are a little brittle and likely deserve using a stronger typing model in the future.
1 parent fe9877a commit 11fdf46

File tree

2 files changed

+153
-6
lines changed

2 files changed

+153
-6
lines changed

databricks/sdk/credentials_provider.py

Lines changed: 8 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -188,19 +188,19 @@ def token() -> Token:
188188
def external_browser(cfg: 'Config') -> Optional[CredentialsProvider]:
189189
if cfg.auth_type != 'external-browser':
190190
return None
191+
191192
client_id, client_secret = None, None
192193
if cfg.client_id:
193194
client_id = cfg.client_id
194195
client_secret = cfg.client_secret
195196
elif cfg.azure_client_id:
196197
client_id = cfg.azure_client
197198
client_secret = cfg.azure_client_secret
198-
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,12 +210,13 @@ 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:
214+
# Pro-actively refresh the loaded credentials. This is done
215+
# to detect if the token is expired and needs to be refreshed
216+
# by going through the OAuth login flow.
216217
credentials.token()
217218
return credentials(cfg)
218-
# TODO: we should ideally use more specific exceptions.
219+
# TODO: We should ideally use more specific exceptions.
219220
except Exception as e:
220221
logger.warning(f'Failed to refresh cached token: {e}. Initiating new OAuth login flow')
221222

@@ -226,6 +227,7 @@ def external_browser(cfg: 'Config') -> Optional[CredentialsProvider]:
226227
consent = oauth_client.initiate_consent()
227228
if not consent:
228229
return None
230+
229231
credentials = consent.launch_external_browser()
230232
token_cache.save(credentials)
231233
return credentials(cfg)

tests/test_credentials_provider.py

Lines changed: 145 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,145 @@
1+
from unittest.mock import Mock
2+
3+
from databricks.sdk.credentials_provider import external_browser
4+
5+
6+
def test_external_browser_refresh_success(mocker):
7+
"""Tests successful refresh of existing credentials."""
8+
9+
# 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+
# 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 SessionCredentials.
24+
want_credentials_provider = lambda c: "new_credentials"
25+
mock_session_credentials.return_value = want_credentials_provider
26+
27+
# Inject the mock implementations.
28+
mocker.patch('databricks.sdk.credentials_provider.TokenCache', return_value=mock_token_cache)
29+
30+
got_credentials_provider = external_browser(mock_cfg)
31+
32+
mock_token_cache.load.assert_called_once()
33+
mock_session_credentials.token.assert_called_once() # Verify token refresh was attempted
34+
assert got_credentials_provider == want_credentials_provider
35+
36+
37+
def test_external_browser_refresh_failure_new_oauth_flow(mocker):
38+
"""Tests failed refresh, triggering a new OAuth flow."""
39+
40+
# Mock Config.
41+
mock_cfg = Mock()
42+
mock_cfg.auth_type = 'external-browser'
43+
mock_cfg.host = 'test-host'
44+
mock_cfg.oidc_endpoints = {'token_endpoint': 'test-token-endpoint'}
45+
mock_cfg.client_id = 'test-client-id'
46+
mock_cfg.client_secret = 'test-client-secret'
47+
48+
# Mock TokenCache.
49+
mock_token_cache = Mock()
50+
mock_session_credentials = Mock()
51+
mock_session_credentials.token.side_effect = Exception(
52+
"Simulated refresh error") # Simulate a failed refresh
53+
mock_token_cache.load.return_value = mock_session_credentials
54+
55+
# Mock SessionCredentials.
56+
want_credentials_provider = lambda c: "new_credentials"
57+
mock_session_credentials.return_value = want_credentials_provider
58+
59+
# Mock OAuthClient.
60+
mock_oauth_client = Mock()
61+
mock_consent = Mock()
62+
mock_consent.launch_external_browser.return_value = mock_session_credentials
63+
mock_oauth_client.initiate_consent.return_value = mock_consent
64+
65+
# Inject the mock implementations.
66+
mocker.patch('databricks.sdk.credentials_provider.TokenCache', return_value=mock_token_cache)
67+
mocker.patch('databricks.sdk.credentials_provider.OAuthClient', return_value=mock_oauth_client)
68+
69+
got_credentials_provider = external_browser(mock_cfg)
70+
71+
mock_token_cache.load.assert_called_once()
72+
mock_session_credentials.token.assert_called_once() # Refresh attempt
73+
mock_oauth_client.initiate_consent.assert_called_once()
74+
mock_consent.launch_external_browser.assert_called_once()
75+
mock_token_cache.save.assert_called_once_with(mock_session_credentials)
76+
assert got_credentials_provider == want_credentials_provider
77+
78+
79+
def test_external_browser_no_cached_credentials(mocker):
80+
"""Tests the case where there are no cached credentials, initiating a new OAuth flow."""
81+
82+
# Mock Config.
83+
mock_cfg = Mock()
84+
mock_cfg.auth_type = 'external-browser'
85+
mock_cfg.host = 'test-host'
86+
mock_cfg.oidc_endpoints = {'token_endpoint': 'test-token-endpoint'}
87+
mock_cfg.client_id = 'test-client-id'
88+
mock_cfg.client_secret = 'test-client-secret'
89+
90+
# Mock TokenCache.
91+
mock_token_cache = Mock()
92+
mock_token_cache.load.return_value = None # No cached credentials
93+
94+
# Mock SessionCredentials.
95+
mock_session_credentials = Mock()
96+
want_credentials_provider = lambda c: "new_credentials"
97+
mock_session_credentials.return_value = want_credentials_provider
98+
99+
# Mock OAuthClient.
100+
mock_consent = Mock()
101+
mock_consent.launch_external_browser.return_value = mock_session_credentials
102+
mock_oauth_client = Mock()
103+
mock_oauth_client.initiate_consent.return_value = mock_consent
104+
105+
# Inject the mock implementations.
106+
mocker.patch('databricks.sdk.credentials_provider.TokenCache', return_value=mock_token_cache)
107+
mocker.patch('databricks.sdk.credentials_provider.OAuthClient', return_value=mock_oauth_client)
108+
109+
got_credentials_provider = external_browser(mock_cfg)
110+
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 got_credentials_provider == want_credentials_provider
116+
117+
118+
def test_external_browser_consent_fails(mocker):
119+
"""Tests the case where OAuth consent initiation fails."""
120+
121+
# 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+
# Mock TokenCache.
130+
mock_token_cache = Mock()
131+
mock_token_cache.load.return_value = None # No cached credentials
132+
133+
# Mock OAuthClient.
134+
mock_oauth_client = Mock()
135+
mock_oauth_client.initiate_consent.return_value = None # Simulate consent failure
136+
137+
# Inject the mock implementations.
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+
got_credentials_provider = external_browser(mock_cfg)
142+
143+
mock_token_cache.load.assert_called_once()
144+
mock_oauth_client.initiate_consent.assert_called_once()
145+
assert got_credentials_provider is None

0 commit comments

Comments
 (0)