Skip to content

Commit 77aa49e

Browse files
jayantsing-dbjayantsing-db
authored
Introduce 'FakeService' to Optimize E2E Integration Tests (#168)
This commit introduces the 'FakeService,' a JUnit 5 extension designed to enhance the end-to-end integration testing process of the OSS JDBC driver. The extension operates in two modes, determined by the 'FAKE_SERVICE_TEST_MODE' environment variable: - RECORD mode: The embedded proxy server intercepts and stores all HTTP requests and their corresponding responses. - REPLAY mode: The embedded proxy server intercepts HTTP requests and returns responses based on the saved stubbings. In REPLAY mode, requests bypass production services, significantly reducing end-to-end test execution time and minimizing cluster costs associated with testing. These optimizations are expected to positively impact OSS CI/CD pipelines. In tests involving 5-6 end-to-end statement-execution scenarios, the runtime has notably improved from 1-2 minutes to approximately 500 milliseconds. However, there is an added cost of 15-20 seconds when running in RECORD mode because of persisting stubbings. * Add StubMappingTransformer to clean credentials from stub mappings * Address review suggestions * Increase unit test coverage * Reset DatabricksHttpClient in integration tests
1 parent a66d7fe commit 77aa49e

File tree

12 files changed

+1357
-45
lines changed

12 files changed

+1357
-45
lines changed

‎pom.xml‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,7 @@
4545
<annotation.version>1.3.5</annotation.version>
4646
<slt.executor>dbsql</slt.executor>
4747
<slt.token>dummy-token</slt.token>
48+
<wiremock.version>3.5.4</wiremock.version>
4849
</properties>
4950
<dependencies>
5051
<dependency>
@@ -168,6 +169,18 @@
168169
<artifactId>jakarta.annotation-api</artifactId>
169170
<version>${annotation.version}</version>
170171
</dependency>
172+
<dependency>
173+
<groupId>org.wiremock</groupId>
174+
<artifactId>wiremock</artifactId>
175+
<version>${wiremock.version}</version>
176+
<scope>test</scope>
177+
<exclusions>
178+
<exclusion>
179+
<groupId>commons-fileupload</groupId>
180+
<artifactId>commons-fileupload</artifactId>
181+
</exclusion>
182+
</exclusions>
183+
</dependency>
171184
</dependencies>
172185

173186
<build>

‎src/main/java/com/databricks/jdbc/client/http/DatabricksHttpClient.java‎

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,8 @@
11
package com.databricks.jdbc.client.http;
22

3+
import static com.databricks.jdbc.driver.DatabricksJdbcConstants.FAKE_SERVICE_URI_PROP_SUFFIX;
4+
import static com.databricks.jdbc.driver.DatabricksJdbcConstants.IS_FAKE_SERVICE_TEST_PROP;
5+
36
import com.databricks.jdbc.client.DatabricksHttpException;
47
import com.databricks.jdbc.client.IDatabricksHttpClient;
58
import com.databricks.jdbc.driver.IDatabricksConnectionContext;
@@ -21,10 +24,13 @@
2124
import org.apache.http.client.methods.CloseableHttpResponse;
2225
import org.apache.http.client.methods.HttpUriRequest;
2326
import org.apache.http.client.protocol.HttpClientContext;
27+
import org.apache.http.conn.UnsupportedSchemeException;
28+
import org.apache.http.conn.routing.HttpRoute;
2429
import org.apache.http.impl.client.BasicCredentialsProvider;
2530
import org.apache.http.impl.client.CloseableHttpClient;
2631
import org.apache.http.impl.client.HttpClientBuilder;
2732
import org.apache.http.impl.client.ProxyAuthenticationStrategy;
33+
import org.apache.http.impl.conn.DefaultSchemePortResolver;
2834
import org.apache.http.impl.conn.PoolingHttpClientConnectionManager;
2935
import org.apache.http.protocol.HttpContext;
3036
import org.slf4j.Logger;
@@ -155,7 +161,10 @@ public void process(HttpResponse httpResponse, HttpContext httpContext)
155161
connectionContext.getUseProxyAuth(),
156162
connectionContext.getProxyUser(),
157163
connectionContext.getProxyPassword());
164+
} else if (Boolean.parseBoolean(System.getProperty(IS_FAKE_SERVICE_TEST_PROP))) {
165+
setFakeServiceRouteInHttpClient(builder);
158166
}
167+
159168
return builder.build();
160169
}
161170

@@ -179,6 +188,27 @@ public static void setProxyDetailsInHttpClient(
179188
}
180189
}
181190

191+
@VisibleForTesting
192+
static void setFakeServiceRouteInHttpClient(HttpClientBuilder builder) {
193+
builder.setRoutePlanner(
194+
(host, request, context) -> {
195+
// Get the fake service URI for the target URI and set it as proxy
196+
final HttpHost proxy =
197+
HttpHost.create(System.getProperty(host.toURI() + FAKE_SERVICE_URI_PROP_SUFFIX));
198+
final HttpHost target;
199+
try {
200+
target =
201+
new HttpHost(
202+
host.getHostName(),
203+
DefaultSchemePortResolver.INSTANCE.resolve(host),
204+
host.getSchemeName());
205+
} catch (UnsupportedSchemeException e) {
206+
throw new HttpException(e.getMessage());
207+
}
208+
return new HttpRoute(target, null, proxy, false);
209+
});
210+
}
211+
182212
@VisibleForTesting
183213
static boolean isRetryAllowed(String method) {
184214
// For now, allowing retry only for GET which is idempotent
@@ -246,4 +276,17 @@ static String getUserAgent() {
246276
}
247277
return mergedString.toString();
248278
}
279+
280+
/** Reset the instance of the http client. This is used for testing purposes only. */
281+
@VisibleForTesting
282+
public static synchronized void resetInstance() {
283+
if (instance != null) {
284+
try {
285+
instance.httpClient.close();
286+
} catch (IOException e) {
287+
LOGGER.error("Caught error while closing http client", e);
288+
}
289+
instance = null;
290+
}
291+
}
249292
}

‎src/main/java/com/databricks/jdbc/core/ArrowResultChunk.java‎

Lines changed: 4 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22

33
import static com.databricks.jdbc.client.impl.thrift.commons.DatabricksThriftHelper.createExternalLink;
44
import static com.databricks.jdbc.commons.util.ValidationUtil.checkHTTPError;
5+
import static com.databricks.jdbc.driver.DatabricksJdbcConstants.IS_FAKE_SERVICE_TEST_PROP;
56

67
import com.databricks.jdbc.client.DatabricksHttpException;
78
import com.databricks.jdbc.client.IDatabricksHttpClient;
@@ -261,7 +262,8 @@ void setStatus(ChunkStatus status) {
261262
/** Checks if the link is valid */
262263
boolean isChunkLinkInvalid() {
263264
return status == ChunkStatus.PENDING
264-
|| expiryTime.minusSeconds(SECONDS_BUFFER_FOR_EXPIRY).isBefore(Instant.now());
265+
|| (!Boolean.parseBoolean(System.getProperty(IS_FAKE_SERVICE_TEST_PROP))
266+
&& expiryTime.minusSeconds(SECONDS_BUFFER_FOR_EXPIRY).isBefore(Instant.now()));
265267
}
266268

267269
/** Returns the status for the chunk */
@@ -396,11 +398,7 @@ synchronized boolean releaseChunk() {
396398
return true;
397399
}
398400

399-
/**
400-
* Returns number of recordBatches in the chunk.
401-
*
402-
* @return
403-
*/
401+
/** Returns number of recordBatches in the chunk. */
404402
int getRecordBatchCountInChunk() {
405403
return this.isDataInitialized ? this.recordBatchList.size() : 0;
406404
}

‎src/main/java/com/databricks/jdbc/driver/DatabricksJdbcConstants.java‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
package com.databricks.jdbc.driver;
22

3+
import com.google.common.annotations.VisibleForTesting;
34
import java.util.Map;
45
import java.util.Set;
56
import java.util.regex.Pattern;
@@ -114,5 +115,16 @@ public final class DatabricksJdbcConstants {
114115
public static final Set<String> ALLOWED_CLIENT_INFO_PROPERTIES =
115116
Set.of(ALLOWED_VOLUME_INGESTION_PATHS);
116117

118+
@VisibleForTesting public static final String IS_FAKE_SERVICE_TEST_PROP = "isFakeServiceTest";
119+
120+
@VisibleForTesting public static final String FAKE_SERVICE_URI_PROP_SUFFIX = ".fakeServiceURI";
121+
122+
/** Enum for the services that can be replaced with a fake service in integration tests. */
123+
@VisibleForTesting
124+
public enum FakeServiceType {
125+
SQL_EXEC,
126+
CLOUD_FETCH
127+
}
128+
117129
public static final String USE_THRIFT_CLIENT = "usethriftclient";
118130
}

‎src/test/java/com/databricks/jdbc/client/http/DatabricksHttpClientTest.java‎

Lines changed: 134 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,9 @@
11
package com.databricks.jdbc.client.http;
22

3-
import static com.databricks.jdbc.client.http.DatabricksHttpClient.*;
3+
import static com.databricks.jdbc.client.http.DatabricksHttpClient.isErrorCodeRetryable;
4+
import static com.databricks.jdbc.client.http.DatabricksHttpClient.isRetryAllowed;
5+
import static com.databricks.jdbc.driver.DatabricksJdbcConstants.FAKE_SERVICE_URI_PROP_SUFFIX;
6+
import static com.databricks.jdbc.driver.DatabricksJdbcConstants.IS_FAKE_SERVICE_TEST_PROP;
47
import static org.junit.jupiter.api.Assertions.*;
58
import static org.mockito.Mockito.*;
69

@@ -9,17 +12,25 @@
912
import com.databricks.jdbc.driver.DatabricksDriver;
1013
import com.databricks.jdbc.driver.IDatabricksConnectionContext;
1114
import java.io.IOException;
15+
import java.lang.reflect.Field;
1216
import java.net.URI;
1317
import java.util.Properties;
1418
import java.util.concurrent.TimeUnit;
19+
import org.apache.http.HttpException;
20+
import org.apache.http.HttpHost;
1521
import org.apache.http.client.methods.CloseableHttpResponse;
22+
import org.apache.http.client.methods.HttpGet;
1623
import org.apache.http.client.methods.HttpUriRequest;
24+
import org.apache.http.conn.routing.HttpRoute;
25+
import org.apache.http.conn.routing.HttpRoutePlanner;
1726
import org.apache.http.impl.client.CloseableHttpClient;
1827
import org.apache.http.impl.client.HttpClientBuilder;
1928
import org.apache.http.impl.conn.PoolingHttpClientConnectionManager;
2029
import org.junit.jupiter.api.Test;
2130
import org.junit.jupiter.api.extension.ExtendWith;
31+
import org.mockito.ArgumentCaptor;
2232
import org.mockito.Mock;
33+
import org.mockito.Mockito;
2334
import org.mockito.junit.jupiter.MockitoExtension;
2435

2536
@ExtendWith(MockitoExtension.class)
@@ -60,6 +71,87 @@ public void testSetProxyDetailsIntoHttpClient() {
6071
builder, null, 8080, true, "user", "proxyPassword"));
6172
}
6273

74+
@Test
75+
public void testSetFakeServiceRouteInHttpClient() throws HttpException {
76+
final String testTargetURI = "https://example.com";
77+
final String testFakeServiceURI = "http://localhost:8080";
78+
System.setProperty(testTargetURI + FAKE_SERVICE_URI_PROP_SUFFIX, testFakeServiceURI);
79+
80+
HttpClientBuilder httpClientBuilder = Mockito.mock(HttpClientBuilder.class);
81+
ArgumentCaptor<HttpRoutePlanner> routePlannerCaptor =
82+
ArgumentCaptor.forClass(HttpRoutePlanner.class);
83+
84+
DatabricksHttpClient.setFakeServiceRouteInHttpClient(httpClientBuilder);
85+
86+
// Capture the route planner set in builder
87+
Mockito.verify(httpClientBuilder).setRoutePlanner(routePlannerCaptor.capture());
88+
HttpRoutePlanner capturedRoutePlanner = routePlannerCaptor.getValue();
89+
90+
// Create a request and determine the route
91+
HttpGet request = new HttpGet(testTargetURI);
92+
HttpHost proxy = HttpHost.create(testFakeServiceURI);
93+
HttpRoute route =
94+
capturedRoutePlanner.determineRoute(
95+
HttpHost.create(request.getURI().toString()), request, null);
96+
97+
// Verify the route is set to the fake service URI
98+
assertEquals(proxy, route.getProxyHost());
99+
assertEquals(2, route.getHopCount());
100+
101+
System.clearProperty(testTargetURI + FAKE_SERVICE_URI_PROP_SUFFIX);
102+
}
103+
104+
@Test
105+
public void testSetFakeServiceRouteInHttpClientThrowsError() {
106+
final String testTargetURI = "https://example.com";
107+
108+
HttpClientBuilder httpClientBuilder = Mockito.mock(HttpClientBuilder.class);
109+
ArgumentCaptor<HttpRoutePlanner> routePlannerCaptor =
110+
ArgumentCaptor.forClass(HttpRoutePlanner.class);
111+
112+
DatabricksHttpClient.setFakeServiceRouteInHttpClient(httpClientBuilder);
113+
114+
Mockito.verify(httpClientBuilder).setRoutePlanner(routePlannerCaptor.capture());
115+
HttpRoutePlanner capturedRoutePlanner = routePlannerCaptor.getValue();
116+
117+
HttpGet request = new HttpGet(testTargetURI);
118+
119+
// Determine route should throw an error as the fake service URI is not set
120+
assertThrows(
121+
IllegalArgumentException.class,
122+
() ->
123+
capturedRoutePlanner.determineRoute(
124+
HttpHost.create(request.getURI().toString()), request, null));
125+
}
126+
127+
@Test
128+
public void testSetFakeServiceRouteInHttpClientThrowsHTTPError() {
129+
// Invalid scheme
130+
final String testTargetURI = "invalid://example.com";
131+
final String testFakeServiceURI = "http://localhost:8080";
132+
System.setProperty(testTargetURI + FAKE_SERVICE_URI_PROP_SUFFIX, testFakeServiceURI);
133+
134+
HttpClientBuilder httpClientBuilder = Mockito.mock(HttpClientBuilder.class);
135+
ArgumentCaptor<HttpRoutePlanner> routePlannerCaptor =
136+
ArgumentCaptor.forClass(HttpRoutePlanner.class);
137+
138+
DatabricksHttpClient.setFakeServiceRouteInHttpClient(httpClientBuilder);
139+
140+
Mockito.verify(httpClientBuilder).setRoutePlanner(routePlannerCaptor.capture());
141+
HttpRoutePlanner capturedRoutePlanner = routePlannerCaptor.getValue();
142+
143+
HttpGet request = new HttpGet(testTargetURI);
144+
145+
// Determine route should throw HTTP error as the target URI is invalid
146+
assertThrows(
147+
HttpException.class,
148+
() ->
149+
capturedRoutePlanner.determineRoute(
150+
HttpHost.create(request.getURI().toString()), request, null));
151+
152+
System.clearProperty(testTargetURI + FAKE_SERVICE_URI_PROP_SUFFIX);
153+
}
154+
63155
@Test
64156
void testExecuteThrowsError() throws IOException {
65157
DatabricksHttpClient databricksHttpClient =
@@ -129,4 +221,45 @@ void testUserAgent() throws Exception {
129221
assertTrue(userAgent.contains(" databricks-jdbc-http "));
130222
assertFalse(userAgent.contains("databricks-sdk-java"));
131223
}
224+
225+
@Test
226+
public void testResetInstance() {
227+
System.setProperty(IS_FAKE_SERVICE_TEST_PROP, "true");
228+
229+
IDatabricksConnectionContext connectionContext =
230+
Mockito.mock(IDatabricksConnectionContext.class);
231+
when(connectionContext.getUseSystemProxy()).thenReturn(false);
232+
when(connectionContext.getUseProxy()).thenReturn(false);
233+
when(connectionContext.getUseCloudFetchProxy()).thenReturn(false);
234+
235+
DatabricksHttpClient databricksHttpClient = DatabricksHttpClient.getInstance(connectionContext);
236+
assertNotNull(databricksHttpClient);
237+
238+
DatabricksHttpClient.resetInstance();
239+
240+
DatabricksHttpClient newInstance = DatabricksHttpClient.getInstance(connectionContext);
241+
assertNotNull(newInstance);
242+
// The instance should be different after reset
243+
assertNotSame(databricksHttpClient, newInstance);
244+
245+
System.clearProperty(IS_FAKE_SERVICE_TEST_PROP);
246+
}
247+
248+
@Test
249+
void testResetInstanceCatchesIOException()
250+
throws IOException, NoSuchFieldException, IllegalAccessException {
251+
DatabricksHttpClient testInstance = new DatabricksHttpClient(mockHttpClient, connectionManager);
252+
253+
doThrow(new IOException()).when(mockHttpClient).close();
254+
255+
// Set the instance to the testInstance using reflection
256+
Field instanceField = DatabricksHttpClient.class.getDeclaredField("instance");
257+
instanceField.setAccessible(true);
258+
instanceField.set(null, testInstance);
259+
260+
DatabricksHttpClient.resetInstance();
261+
262+
verify(mockHttpClient, times(1)).close();
263+
assertNull(instanceField.get(null));
264+
}
132265
}

0 commit comments

Comments
 (0)