Skip to content

Commit 54ffcc4

Browse files
Tejas-KocharTejas-Kochar
committed
tests for custom scopes support, scopes sorting and scopes parsing
1 parent 724548f commit 54ffcc4

File tree

7 files changed

+459
-9
lines changed

7 files changed

+459
-9
lines changed

‎databricks-sdk-java/src/main/java/com/databricks/sdk/core/ConfigAttributeAccessor.java‎

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,9 +2,13 @@
22

33
import com.databricks.sdk.support.InternalApi;
44
import java.lang.reflect.Field;
5+
import java.lang.reflect.ParameterizedType;
56
import java.time.Duration;
7+
import java.util.Arrays;
8+
import java.util.List;
69
import java.util.Map;
710
import java.util.Objects;
11+
import java.util.stream.Collectors;
812

913
@InternalApi
1014
class ConfigAttributeAccessor {
@@ -63,6 +67,21 @@ public void setValueOnConfig(DatabricksConfig cfg, String value) throws IllegalA
6367
field.set(cfg, seconds > 0 ? Duration.ofSeconds(seconds) : null);
6468
} else if (field.getType() == ProxyConfig.ProxyAuthType.class) {
6569
field.set(cfg, ProxyConfig.ProxyAuthType.valueOf(value));
70+
} else if (List.class.isAssignableFrom(field.getType())) {
71+
// Handle List<String> fields (e.g., scopes)
72+
// Parse comma and/or whitespace separated values from environment variable or config file
73+
if (field.getGenericType() instanceof ParameterizedType) {
74+
ParameterizedType paramType = (ParameterizedType) field.getGenericType();
75+
if (paramType.getActualTypeArguments().length > 0
76+
&& paramType.getActualTypeArguments()[0] == String.class) {
77+
// Split by commas and/or whitespace and filter out empty strings
78+
List<String> list =
79+
Arrays.stream(value.trim().split("[,\\s]+"))
80+
.filter(s -> !s.isEmpty())
81+
.collect(Collectors.toList());
82+
field.set(cfg, list);
83+
}
84+
}
6685
}
6786
field.setAccessible(false);
6887
}
@@ -91,6 +110,10 @@ public String toString() {
91110
}
92111

93112
public String getAsString(Object value) {
113+
if (value instanceof List) {
114+
// Format lists as space-separated values for consistency with input format
115+
return String.join(" ", (List<String>) value);
116+
}
94117
return value.toString();
95118
}
96119

‎databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -204,6 +204,7 @@ private synchronized DatabricksConfig innerResolve() {
204204
try {
205205
ConfigLoader.resolve(this);
206206
ConfigLoader.validate(this);
207+
sortScopes();
207208
ConfigLoader.fixHostIfNeeded(this);
208209
initHttp();
209210
return this;
@@ -212,6 +213,15 @@ private synchronized DatabricksConfig innerResolve() {
212213
}
213214
}
214215

216+
/**
217+
* Sort scopes in-place for better de-duplication in the refresh token cache.
218+
*/
219+
private void sortScopes() {
220+
if (scopes != null && !scopes.isEmpty()) {
221+
java.util.Collections.sort(scopes);
222+
}
223+
}
224+
215225
private void initHttp() {
216226
if (httpClient != null) {
217227
return;

‎databricks-sdk-java/src/test/java/com/databricks/sdk/core/DatabricksConfigTest.java‎

Lines changed: 48 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -322,4 +322,52 @@ public void testDisableOauthRefreshTokenEnvironmentVariable() {
322322

323323
assertEquals(true, config.getDisableOauthRefreshToken());
324324
}
325+
326+
// Config File Scope Parsing Tests
327+
328+
@Test
329+
public void testConfigFileScopesEmptyDefaultsToAllApis() {
330+
Map<String, String> env = new HashMap<>();
331+
env.put("HOME", "src/test/resources/testdata");
332+
333+
DatabricksConfig config = new DatabricksConfig().setProfile("scope-empty");
334+
config.resolve(new Environment(env, new ArrayList<>(), System.getProperty("os.name")));
335+
336+
List<String> scopes = config.getScopes();
337+
assertEquals(1, scopes.size());
338+
assertEquals("all-apis", scopes.get(0));
339+
}
340+
341+
@Test
342+
public void testConfigFileScopesSingle() {
343+
Map<String, String> env = new HashMap<>();
344+
env.put("HOME", "src/test/resources/testdata");
345+
346+
DatabricksConfig config = new DatabricksConfig().setProfile("scope-single");
347+
config.resolve(new Environment(env, new ArrayList<>(), System.getProperty("os.name")));
348+
349+
List<String> scopes = config.getScopes();
350+
assertEquals(1, scopes.size());
351+
assertEquals("clusters:read", scopes.get(0));
352+
}
353+
354+
@Test
355+
public void testConfigFileScopesMultipleSorted() {
356+
Map<String, String> env = new HashMap<>();
357+
env.put("HOME", "src/test/resources/testdata");
358+
359+
DatabricksConfig config = new DatabricksConfig().setProfile("scope-multiple");
360+
config.resolve(new Environment(env, new ArrayList<>(), System.getProperty("os.name")));
361+
362+
List<String> scopes = config.getScopes();
363+
// Should be sorted alphabetically
364+
assertEquals(7, scopes.size());
365+
assertEquals("clusters:read", scopes.get(0));
366+
assertEquals("files:read", scopes.get(1));
367+
assertEquals("iam:read", scopes.get(2));
368+
assertEquals("jobs:read", scopes.get(3));
369+
assertEquals("mlflow:read", scopes.get(4));
370+
assertEquals("model-serving:read", scopes.get(5));
371+
assertEquals("pipelines:read", scopes.get(6));
372+
}
325373
}

‎databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/DatabricksOAuthTokenSourceTest.java‎

Lines changed: 113 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,9 +12,12 @@
1212
import java.io.IOException;
1313
import java.net.MalformedURLException;
1414
import java.net.URL;
15+
import java.util.Arrays;
1516
import java.util.HashMap;
17+
import java.util.List;
1618
import java.util.Map;
1719
import java.util.stream.Stream;
20+
import org.junit.jupiter.api.Test;
1821
import org.junit.jupiter.params.ParameterizedTest;
1922
import org.junit.jupiter.params.provider.MethodSource;
2023
import org.mockito.Mockito;
@@ -335,4 +338,114 @@ void testTokenSource(TestCase testCase) {
335338
verify(testCase.idTokenSource, atLeastOnce()).getIDToken(testCase.expectedAudience);
336339
}
337340
}
341+
342+
// Scope-specific tests for WIF/OIDC token exchange
343+
344+
@Test
345+
void testDefaultScopesInTokenExchange() throws IOException {
346+
// Verify default "all-apis" scope is used when no scopes specified
347+
OpenIDConnectEndpoints testEndpoints =
348+
new OpenIDConnectEndpoints(TEST_TOKEN_ENDPOINT, TEST_AUTHORIZATION_ENDPOINT);
349+
IDTokenSource testIdTokenSource = Mockito.mock(IDTokenSource.class);
350+
when(testIdTokenSource.getIDToken(any())).thenReturn(new IDToken(TEST_ID_TOKEN));
351+
352+
Map<String, Object> successResponse = new HashMap<>();
353+
successResponse.put("access_token", TOKEN);
354+
successResponse.put("token_type", TOKEN_TYPE);
355+
successResponse.put("expires_in", EXPIRES_IN);
356+
String successJson = new ObjectMapper().writeValueAsString(successResponse);
357+
358+
// Expected request with default scope
359+
Map<String, String> formParams = new HashMap<>();
360+
formParams.put("client_id", TEST_CLIENT_ID);
361+
formParams.put("subject_token", TEST_ID_TOKEN);
362+
formParams.put("subject_token_type", "urn:ietf:params:oauth:token-type:jwt");
363+
formParams.put("grant_type", "urn:ietf:params:oauth:grant-type:token-exchange");
364+
formParams.put("scope", "all-apis");
365+
FormRequest expectedRequest = new FormRequest(TEST_TOKEN_ENDPOINT, formParams);
366+
367+
HttpClient mockHttpClient = createMockHttpClient(expectedRequest, 200, successJson);
368+
369+
DatabricksOAuthTokenSource tokenSource =
370+
new DatabricksOAuthTokenSource.Builder(
371+
TEST_CLIENT_ID, TEST_HOST, testEndpoints, testIdTokenSource, mockHttpClient)
372+
.scopes(Arrays.asList("all-apis"))
373+
.build();
374+
375+
Token token = tokenSource.getToken();
376+
assertEquals(TOKEN, token.getAccessToken());
377+
}
378+
379+
@Test
380+
void testCustomScopesInTokenExchange() throws IOException {
381+
// Verify custom scopes are correctly passed to token exchange
382+
OpenIDConnectEndpoints testEndpoints =
383+
new OpenIDConnectEndpoints(TEST_TOKEN_ENDPOINT, TEST_AUTHORIZATION_ENDPOINT);
384+
IDTokenSource testIdTokenSource = Mockito.mock(IDTokenSource.class);
385+
when(testIdTokenSource.getIDToken(any())).thenReturn(new IDToken(TEST_ID_TOKEN));
386+
387+
Map<String, Object> successResponse = new HashMap<>();
388+
successResponse.put("access_token", TOKEN);
389+
successResponse.put("token_type", TOKEN_TYPE);
390+
successResponse.put("expires_in", EXPIRES_IN);
391+
String successJson = new ObjectMapper().writeValueAsString(successResponse);
392+
393+
// Expected request with custom scope
394+
Map<String, String> formParams = new HashMap<>();
395+
formParams.put("client_id", TEST_CLIENT_ID);
396+
formParams.put("subject_token", TEST_ID_TOKEN);
397+
formParams.put("subject_token_type", "urn:ietf:params:oauth:token-type:jwt");
398+
formParams.put("grant_type", "urn:ietf:params:oauth:grant-type:token-exchange");
399+
formParams.put("scope", "unity-catalog:read");
400+
FormRequest expectedRequest = new FormRequest(TEST_TOKEN_ENDPOINT, formParams);
401+
402+
HttpClient mockHttpClient = createMockHttpClient(expectedRequest, 200, successJson);
403+
404+
DatabricksOAuthTokenSource tokenSource =
405+
new DatabricksOAuthTokenSource.Builder(
406+
TEST_CLIENT_ID, TEST_HOST, testEndpoints, testIdTokenSource, mockHttpClient)
407+
.scopes(Arrays.asList("unity-catalog:read"))
408+
.build();
409+
410+
Token token = tokenSource.getToken();
411+
assertEquals(TOKEN, token.getAccessToken());
412+
}
413+
414+
@Test
415+
void testMultipleScopesJoinedWithSpaces() throws IOException {
416+
// Verify multiple scopes are joined with spaces in token exchange request
417+
OpenIDConnectEndpoints testEndpoints =
418+
new OpenIDConnectEndpoints(TEST_TOKEN_ENDPOINT, TEST_AUTHORIZATION_ENDPOINT);
419+
IDTokenSource testIdTokenSource = Mockito.mock(IDTokenSource.class);
420+
when(testIdTokenSource.getIDToken(any())).thenReturn(new IDToken(TEST_ID_TOKEN));
421+
422+
Map<String, Object> successResponse = new HashMap<>();
423+
successResponse.put("access_token", TOKEN);
424+
successResponse.put("token_type", TOKEN_TYPE);
425+
successResponse.put("expires_in", EXPIRES_IN);
426+
String successJson = new ObjectMapper().writeValueAsString(successResponse);
427+
428+
// Scopes should be sorted before being passed (sorted in DatabricksConfig.resolve())
429+
List<String> scopes = Arrays.asList("clusters:read", "jobs:read", "mlflow:read");
430+
431+
// Expected request with multiple scopes joined by spaces
432+
Map<String, String> formParams = new HashMap<>();
433+
formParams.put("client_id", TEST_CLIENT_ID);
434+
formParams.put("subject_token", TEST_ID_TOKEN);
435+
formParams.put("subject_token_type", "urn:ietf:params:oauth:token-type:jwt");
436+
formParams.put("grant_type", "urn:ietf:params:oauth:grant-type:token-exchange");
437+
formParams.put("scope", "clusters:read jobs:read mlflow:read");
438+
FormRequest expectedRequest = new FormRequest(TEST_TOKEN_ENDPOINT, formParams);
439+
440+
HttpClient mockHttpClient = createMockHttpClient(expectedRequest, 200, successJson);
441+
442+
DatabricksOAuthTokenSource tokenSource =
443+
new DatabricksOAuthTokenSource.Builder(
444+
TEST_CLIENT_ID, TEST_HOST, testEndpoints, testIdTokenSource, mockHttpClient)
445+
.scopes(scopes)
446+
.build();
447+
448+
Token token = tokenSource.getToken();
449+
assertEquals(TOKEN, token.getAccessToken());
450+
}
338451
}

0 commit comments

Comments
 (0)