Skip to content

Commit 01a0bc2

Browse files
emmyzhou-dbemmyzhou-db
committed
Update unit tests
1 parent 71a8d3d commit 01a0bc2

File tree

2 files changed

+136
-120
lines changed

2 files changed

+136
-120
lines changed

‎databricks-sdk-java/src/test/java/com/databricks/sdk/DatabricksAuthTest.java‎

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -4,22 +4,26 @@
44
package com.databricks.sdk;
55

66
import static org.junit.jupiter.api.Assertions.*;
7+
import static org.mockito.Mockito.mock;
78

89
import com.databricks.sdk.core.ConfigResolving;
910
import com.databricks.sdk.core.DatabricksConfig;
10-
import com.databricks.sdk.core.DummyHttpClient;
11+
import com.databricks.sdk.core.http.HttpClient;
1112
import com.databricks.sdk.core.utils.GitHubUtils;
1213
import com.databricks.sdk.core.utils.TestOSUtils;
1314
import java.io.File;
1415
import org.junit.jupiter.api.Test;
16+
import org.junit.jupiter.api.extension.ExtendWith;
17+
import org.mockito.junit.jupiter.MockitoExtension;
1518

19+
@ExtendWith(MockitoExtension.class)
1620
public class DatabricksAuthTest implements GitHubUtils, ConfigResolving {
1721

1822
private String errorMessageBase =
1923
"default auth: cannot configure default credentials, please check https://docs.databricks.com/en/dev-tools/auth.html#databricks-client-unified-authentication to configure credentials for your preferred authentication method";
2024

2125
private DatabricksConfig createConfigWithMockClient() {
22-
return new DatabricksConfig().setHttpClient(new DummyHttpClient());
26+
return new DatabricksConfig().setHttpClient(mock(HttpClient.class));
2327
}
2428

2529
public DatabricksAuthTest() {

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

Lines changed: 130 additions & 118 deletions
Original file line numberDiff line numberDiff line change
@@ -8,144 +8,156 @@
88
import com.databricks.sdk.core.http.HttpClient;
99
import com.databricks.sdk.core.http.Request;
1010
import com.databricks.sdk.core.http.Response;
11-
import com.fasterxml.jackson.databind.ObjectMapper;
12-
import com.fasterxml.jackson.databind.node.ObjectNode;
1311
import java.io.IOException;
1412
import java.net.MalformedURLException;
1513
import java.net.URL;
1614
import java.util.stream.Stream;
17-
import org.junit.jupiter.api.BeforeEach;
18-
import org.junit.jupiter.api.Test;
1915
import org.junit.jupiter.params.ParameterizedTest;
2016
import org.junit.jupiter.params.provider.Arguments;
2117
import org.junit.jupiter.params.provider.MethodSource;
22-
import org.mockito.Mock;
23-
import org.mockito.MockitoAnnotations;
2418

2519
public class GithubIDTokenSourceTest {
2620
private static final String TEST_REQUEST_URL = "https://github.com/token";
2721
private static final String TEST_REQUEST_TOKEN = "test-request-token";
2822
private static final String TEST_ID_TOKEN = "test-id-token";
2923
private static final String TEST_AUDIENCE = "test-audience";
3024

31-
@Mock private static HttpClient mockHttpClient;
32-
33-
private GithubIDTokenSource tokenSource;
34-
private ObjectMapper mapper;
35-
36-
@BeforeEach
37-
void setUp() throws IOException {
38-
MockitoAnnotations.openMocks(this);
39-
mapper = new ObjectMapper();
40-
tokenSource = new GithubIDTokenSource(TEST_REQUEST_URL, TEST_REQUEST_TOKEN, mockHttpClient);
41-
}
42-
43-
@Test
44-
void testSuccessfulTokenRetrieval() throws IOException {
45-
// Prepare mock response
46-
ObjectNode responseJson = mapper.createObjectNode();
47-
responseJson.put("value", TEST_ID_TOKEN);
48-
Response mockResponse = makeResponse(responseJson.toString(), 200);
49-
when(mockHttpClient.execute(any(Request.class))).thenReturn(mockResponse);
50-
51-
// Test token retrieval
52-
IDToken token = tokenSource.getIDToken(TEST_AUDIENCE);
53-
54-
assertNotNull(token);
55-
assertEquals(TEST_ID_TOKEN, token.getValue());
56-
57-
// Verify the request was made with correct parameters
58-
verify(mockHttpClient)
59-
.execute(
60-
argThat(
61-
request -> {
62-
return request.getMethod().equals("GET")
63-
&& request.getUrl().startsWith(TEST_REQUEST_URL)
64-
&& request.getUrl().contains("audience=" + TEST_AUDIENCE)
65-
&& request
66-
.getHeaders()
67-
.get("Authorization")
68-
.equals("Bearer " + TEST_REQUEST_TOKEN);
69-
}));
70-
}
71-
72-
@Test
73-
void testSuccessfulTokenRetrievalWithoutAudience() throws IOException {
74-
// Prepare mock response
75-
ObjectNode responseJson = mapper.createObjectNode();
76-
responseJson.put("value", TEST_ID_TOKEN);
77-
Response mockResponse = makeResponse(responseJson.toString(), 200);
78-
when(mockHttpClient.execute(any(Request.class))).thenReturn(mockResponse);
79-
80-
// Test token retrieval without audience
81-
IDToken token = tokenSource.getIDToken("");
82-
83-
assertNotNull(token);
84-
assertEquals(TEST_ID_TOKEN, token.getValue());
85-
86-
// Verify the request was made with correct parameters
87-
verify(mockHttpClient)
88-
.execute(
89-
argThat(
90-
request -> {
91-
return request.getMethod().equals("GET")
92-
&& request.getUrl().equals(TEST_REQUEST_URL)
93-
&& request
94-
.getHeaders()
95-
.get("Authorization")
96-
.equals("Bearer " + TEST_REQUEST_TOKEN);
97-
}));
25+
private static HttpClient createHttpMock(
26+
String responseBody, int statusCode, IOException exception) throws IOException {
27+
HttpClient client = mock(HttpClient.class);
28+
if (exception != null) {
29+
when(client.execute(any(Request.class))).thenThrow(exception);
30+
} else {
31+
when(client.execute(any(Request.class))).thenReturn(makeResponse(responseBody, statusCode));
32+
}
33+
return client;
9834
}
9935

100-
private static Stream<Arguments> provideInvalidConstructorParameters() {
101-
return Stream.of(
102-
Arguments.of("Missing Request URL", null, TEST_REQUEST_TOKEN, mockHttpClient),
103-
Arguments.of("Missing Request Token", TEST_REQUEST_URL, null, mockHttpClient),
104-
Arguments.of("Null HttpClient", TEST_REQUEST_URL, TEST_REQUEST_TOKEN, null));
105-
}
106-
107-
@ParameterizedTest(name = "{0}")
108-
@MethodSource("provideInvalidConstructorParameters")
109-
void testInvalidConstructorParameters(
110-
String testName, String requestUrl, String requestToken, HttpClient httpClient) {
111-
GithubIDTokenSource invalidSource =
112-
new GithubIDTokenSource(requestUrl, requestToken, httpClient);
113-
assertThrows(DatabricksException.class, () -> invalidSource.getIDToken(TEST_AUDIENCE));
114-
}
115-
116-
private static Stream<Arguments> provideHttpErrorScenarios() throws IOException {
117-
HttpClient httpClientError = mock(HttpClient.class);
118-
when(httpClientError.execute(any(Request.class))).thenThrow(new IOException("Network error"));
119-
120-
HttpClient nonSuccessClient = mock(HttpClient.class);
121-
when(nonSuccessClient.execute(any(Request.class)))
122-
.thenReturn(makeResponse("Error response", 400));
123-
124-
HttpClient invalidJsonClient = mock(HttpClient.class);
125-
when(invalidJsonClient.execute(any(Request.class)))
126-
.thenReturn(makeResponse("Invalid json", 200));
127-
128-
HttpClient missingTokenClient = mock(HttpClient.class);
129-
when(missingTokenClient.execute(any(Request.class))).thenReturn(makeResponse("{}", 200));
130-
131-
HttpClient emptyTokenClient = mock(HttpClient.class);
132-
when(emptyTokenClient.execute(any(Request.class)))
133-
.thenReturn(makeResponse("{\"value\":\"\"}", 200));
134-
36+
private static Stream<Arguments> provideAllTestScenarios() throws IOException {
13537
return Stream.of(
136-
Arguments.of("HTTP Client Error", httpClientError),
137-
Arguments.of("Non-Success Status Code", nonSuccessClient),
138-
Arguments.of("Invalid JSON Response", invalidJsonClient),
139-
Arguments.of("Missing Token Value", missingTokenClient),
140-
Arguments.of("Empty Token Value", emptyTokenClient));
38+
// Successful token retrieval with audience
39+
Arguments.of(
40+
"Successful token retrieval with audience",
41+
TEST_REQUEST_URL,
42+
TEST_REQUEST_TOKEN,
43+
createHttpMock("{\"value\":\"" + TEST_ID_TOKEN + "\"}", 200, null),
44+
TEST_AUDIENCE,
45+
(java.util.function.Predicate<Request>)
46+
request ->
47+
request.getMethod().equals("GET")
48+
&& request.getUrl().startsWith(TEST_REQUEST_URL)
49+
&& request.getUrl().contains("audience=" + TEST_AUDIENCE)
50+
&& request
51+
.getHeaders()
52+
.get("Authorization")
53+
.equals("Bearer " + TEST_REQUEST_TOKEN),
54+
null),
55+
// Successful token retrieval without audience
56+
Arguments.of(
57+
"Successful token retrieval without audience",
58+
TEST_REQUEST_URL,
59+
TEST_REQUEST_TOKEN,
60+
createHttpMock("{\"value\":\"" + TEST_ID_TOKEN + "\"}", 200, null),
61+
"",
62+
(java.util.function.Predicate<Request>)
63+
request ->
64+
request.getMethod().equals("GET")
65+
&& request.getUrl().equals(TEST_REQUEST_URL)
66+
&& request
67+
.getHeaders()
68+
.get("Authorization")
69+
.equals("Bearer " + TEST_REQUEST_TOKEN),
70+
null),
71+
// Invalid constructor parameters
72+
Arguments.of(
73+
"Missing Request URL",
74+
null,
75+
TEST_REQUEST_TOKEN,
76+
createHttpMock("{\"value\":\"" + TEST_ID_TOKEN + "\"}", 200, null),
77+
TEST_AUDIENCE,
78+
null,
79+
DatabricksException.class),
80+
Arguments.of(
81+
"Missing Request Token",
82+
TEST_REQUEST_URL,
83+
null,
84+
createHttpMock("{\"value\":\"" + TEST_ID_TOKEN + "\"}", 200, null),
85+
TEST_AUDIENCE,
86+
null,
87+
DatabricksException.class),
88+
Arguments.of(
89+
"Null HttpClient",
90+
TEST_REQUEST_URL,
91+
TEST_REQUEST_TOKEN,
92+
null,
93+
TEST_AUDIENCE,
94+
null,
95+
DatabricksException.class),
96+
// HTTP error scenarios
97+
Arguments.of(
98+
"HTTP Client Error",
99+
TEST_REQUEST_URL,
100+
TEST_REQUEST_TOKEN,
101+
createHttpMock(null, 0, new IOException("Network error")),
102+
TEST_AUDIENCE,
103+
null,
104+
DatabricksException.class),
105+
Arguments.of(
106+
"Non-Success Status Code",
107+
TEST_REQUEST_URL,
108+
TEST_REQUEST_TOKEN,
109+
createHttpMock("{\"error\":\"Bad Request\"}", 400, null),
110+
TEST_AUDIENCE,
111+
null,
112+
DatabricksException.class),
113+
Arguments.of(
114+
"Invalid JSON Response",
115+
TEST_REQUEST_URL,
116+
TEST_REQUEST_TOKEN,
117+
createHttpMock("{invalid json}", 200, null),
118+
TEST_AUDIENCE,
119+
null,
120+
DatabricksException.class),
121+
Arguments.of(
122+
"Missing Token Value",
123+
TEST_REQUEST_URL,
124+
TEST_REQUEST_TOKEN,
125+
createHttpMock("{\"other\":\"field\"}", 200, null),
126+
TEST_AUDIENCE,
127+
null,
128+
DatabricksException.class),
129+
Arguments.of(
130+
"Empty Token Value",
131+
TEST_REQUEST_URL,
132+
TEST_REQUEST_TOKEN,
133+
createHttpMock("{\"value\":\"\"}", 200, null),
134+
TEST_AUDIENCE,
135+
null,
136+
DatabricksException.class));
141137
}
142138

143139
@ParameterizedTest(name = "{0}")
144-
@MethodSource("provideHttpErrorScenarios")
145-
void testHttpErrorScenarios(String testName, HttpClient httpClient) {
146-
GithubIDTokenSource source =
147-
new GithubIDTokenSource(TEST_REQUEST_URL, TEST_REQUEST_TOKEN, httpClient);
148-
assertThrows(DatabricksException.class, () -> source.getIDToken(TEST_AUDIENCE));
140+
@MethodSource("provideAllTestScenarios")
141+
void testAllScenarios(
142+
String testName,
143+
String requestUrl,
144+
String requestToken,
145+
HttpClient httpClient,
146+
String audience,
147+
java.util.function.Predicate<Request> requestValidator,
148+
Class<? extends Exception> expectedException)
149+
throws IOException {
150+
151+
GithubIDTokenSource tokenSource = new GithubIDTokenSource(requestUrl, requestToken, httpClient);
152+
153+
if (expectedException != null) {
154+
assertThrows(expectedException, () -> tokenSource.getIDToken(audience));
155+
} else {
156+
IDToken token = tokenSource.getIDToken(audience);
157+
assertNotNull(token);
158+
assertEquals(TEST_ID_TOKEN, token.getValue());
159+
verify(httpClient).execute(argThat(request -> requestValidator.test(request)));
160+
}
149161
}
150162

151163
private static Response makeResponse(String body, int status) throws MalformedURLException {

0 commit comments

Comments
 (0)