Modify MultiChunkExecutionIntegrationTests to extend AbstractFakeServiceIntegrationTests (#195)

AbstractFakeServiceIntegrationTests abstracts out the injection of fake services for
sql-exec and cloud-fetch. Integration tests using the production sql-exec and
cloud-fetch services can extend the abstract class to automatically use fake services
(RECORD or REPLAY mode based on environment `FAKE_SERVICE_TEST_MODE`). The abstract
class also provides hooks to configure the target URL of sql-exec and cloud-fetch (used
in RECORD mode to capture requests hitting target URL).

Author: jayantsing-db

…ice/BaseFakeServiceIntegrationTests.java …AbstractFakeServiceIntegrationTests.javasrc/test/java/com/databricks/jdbc/integration/fakeservice/BaseFakeServiceIntegrationTests.java renamed to src/test/java/com/databricks/jdbc/integration/fakeservice/AbstractFakeServiceIntegrationTests.java src/test/java/com/databricks/jdbc/integration/fakeservice/BaseFakeServiceIntegrationTests.java renamed to src/test/java/com/databricks/jdbc/integration/fakeservice/AbstractFakeServiceIntegrationTests.java

@@ -12,7 +12,7 @@
  * com.databricks.jdbc.driver.DatabricksJdbcConstants.FakeServiceType#SQL_EXEC} and {@link
  * com.databricks.jdbc.driver.DatabricksJdbcConstants.FakeServiceType#CLOUD_FETCH}.
  */
-public abstract class BaseFakeServiceIntegrationTests {
+public abstract class AbstractFakeServiceIntegrationTests {
 
   /**
    * {@link FakeServiceExtension} for {@link DatabricksJdbcConstants.FakeServiceType#SQL_EXEC}.
@@ -40,13 +40,24 @@ public abstract class BaseFakeServiceIntegrationTests {
           DatabricksJdbcConstants.FakeServiceType.CLOUD_FETCH,
           "https://dbstoragepzjc6kojqibtg.blob.core.windows.net");
 
-  /** Resets the possible mutations for the extensions after all tests have run. */
+  /**
+   * Resets the potential mutations (e.g., URLs set by {@link #setSqlExecApiTargetUrl}, {@link
+   * #setCloudFetchApiTargetUrl}) to meaningful defaults, after all tests have completed.
+   */
   @AfterAll
-  protected static void resetPossibleMutations() {
+  static void resetPossibleMutations() {
     sqlExecApiExtension.setTargetBaseUrl("https://" + System.getenv("DATABRICKS_HOST"));
     cloudFetchApiExtension.setTargetBaseUrl("https://dbstoragepzjc6kojqibtg.blob.core.windows.net");
   }
 
+  protected static void setSqlExecApiTargetUrl(final String sqlExecApiTargetUrl) {
+    sqlExecApiExtension.setTargetBaseUrl(sqlExecApiTargetUrl);
+  }
+
+  protected static void setCloudFetchApiTargetUrl(final String cloudFetchApiTargetUrl) {
+    cloudFetchApiExtension.setTargetBaseUrl(cloudFetchApiTargetUrl);
+  }
+
   protected FakeServiceExtension getSqlExecApiExtension() {
     return sqlExecApiExtension;
   }

src/test/java/com/databricks/jdbc/integration/fakeservice/FakeServiceExtension.java

@@ -104,6 +104,9 @@ public class FakeServiceExtension extends DatabricksWireMockExtension {
   /** Fake service to manage. */
   private final FakeServiceType fakeServiceType;
 
+  /** HTTP port of the WireMock server. */
+  private int wireMockServerHttpPort;
+
   /** Base URL of the target production service. */
   private String targetBaseUrl;
 
@@ -145,8 +148,9 @@ protected void onBeforeAll(WireMockRuntimeInfo wireMockRuntimeInfo, ExtensionCon
         fakeServiceModeValue != null
             ? FakeServiceMode.valueOf(fakeServiceModeValue.toUpperCase())
             : FakeServiceMode.REPLAY;
+    wireMockServerHttpPort = wireMockRuntimeInfo.getHttpPort();
 
-    setFakeServiceProperties(wireMockRuntimeInfo);
+    setFakeServiceProperties(wireMockServerHttpPort);
     IntegrationTestUtil.resetJDBCConnection();
     DatabricksHttpClient.resetInstance();
   }
@@ -190,6 +194,9 @@ protected void onAfterAll(WireMockRuntimeInfo wireMockRuntimeInfo, ExtensionCont
   /** Sets the base URL of the target production service to be tracked. */
   protected void setTargetBaseUrl(String targetBaseUrl) {
     this.targetBaseUrl = targetBaseUrl;
+
+    // Refresh the fake service properties
+    setFakeServiceProperties(wireMockServerHttpPort);
   }
 
   /** Gets the stubbing directory for the current test class and method. */
@@ -254,13 +261,12 @@ private void saveStubMappings(WireMockRuntimeInfo wireMockRuntimeInfo, Extension
     new JsonFileMappingsSource(new SingleRootFileSource(stubbingDir), null).save(stubMappingList);
   }
 
-  private void setFakeServiceProperties(WireMockRuntimeInfo wireMockRuntimeInfo) {
+  private void setFakeServiceProperties(int wireMockServerHttpPort) {
     System.setProperty(IS_FAKE_SERVICE_TEST_PROP, "true");
     System.setProperty(
         fakeServiceType.name().toLowerCase() + TARGET_URI_PROP_SUFFIX, targetBaseUrl);
     System.setProperty(
-        targetBaseUrl + FAKE_SERVICE_URI_PROP_SUFFIX,
-        "http://localhost:" + wireMockRuntimeInfo.getHttpPort());
+        targetBaseUrl + FAKE_SERVICE_URI_PROP_SUFFIX, "http://localhost:" + wireMockServerHttpPort);
   }
 
   private void clearFakeServiceProperties() {

src/test/java/com/databricks/jdbc/integration/fakeservice/tests/ConcurrencyIntegrationTests.java

@@ -4,7 +4,7 @@
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.fail;
 
-import com.databricks.jdbc.integration.fakeservice.BaseFakeServiceIntegrationTests;
+import com.databricks.jdbc.integration.fakeservice.AbstractFakeServiceIntegrationTests;
 import java.sql.Connection;
 import java.sql.PreparedStatement;
 import java.sql.ResultSet;
@@ -18,7 +18,7 @@
 import org.junit.jupiter.api.Test;
 
 /** Integration tests for concurrent operations on a table. */
-public class ConcurrencyIntegrationTests extends BaseFakeServiceIntegrationTests {
+public class ConcurrencyIntegrationTests extends AbstractFakeServiceIntegrationTests {
 
   private Connection connection;
 

src/test/java/com/databricks/jdbc/integration/fakeservice/tests/ConnectionIntegrationTests.java

@@ -4,14 +4,14 @@
 import static org.junit.jupiter.api.Assertions.assertThrows;
 
 import com.databricks.jdbc.core.DatabricksSQLException;
-import com.databricks.jdbc.integration.fakeservice.BaseFakeServiceIntegrationTests;
+import com.databricks.jdbc.integration.fakeservice.AbstractFakeServiceIntegrationTests;
 import java.sql.Connection;
 import java.sql.DriverManager;
 import java.sql.SQLException;
 import org.junit.jupiter.api.Test;
 
 /** Integration tests for connection to Databricks service. */
-public class ConnectionIntegrationTests extends BaseFakeServiceIntegrationTests {
+public class ConnectionIntegrationTests extends AbstractFakeServiceIntegrationTests {
 
   @Test
   void testSuccessfulConnection() throws SQLException {

src/test/java/com/databricks/jdbc/integration/fakeservice/tests/ErrorHandlingIntegrationTests.java

@@ -6,12 +6,12 @@
 import com.databricks.jdbc.core.DatabricksParsingException;
 import com.databricks.jdbc.core.DatabricksSQLException;
 import com.databricks.jdbc.core.DatabricksSQLFeatureNotSupportedException;
-import com.databricks.jdbc.integration.fakeservice.BaseFakeServiceIntegrationTests;
+import com.databricks.jdbc.integration.fakeservice.AbstractFakeServiceIntegrationTests;
 import java.sql.*;
 import org.junit.jupiter.api.Test;
 
 /** Integration tests for error handling scenarios. */
-public class ErrorHandlingIntegrationTests extends BaseFakeServiceIntegrationTests {
+public class ErrorHandlingIntegrationTests extends AbstractFakeServiceIntegrationTests {
 
   @Test
   void testFailureToLoadDriver() {

src/test/java/com/databricks/jdbc/integration/fakeservice/tests/ExecutionIntegrationTests.java

@@ -7,14 +7,14 @@
 import static com.github.tomakehurst.wiremock.client.WireMock.urlEqualTo;
 import static org.junit.jupiter.api.Assertions.*;
 
-import com.databricks.jdbc.integration.fakeservice.BaseFakeServiceIntegrationTests;
+import com.databricks.jdbc.integration.fakeservice.AbstractFakeServiceIntegrationTests;
 import com.github.tomakehurst.wiremock.client.CountMatchingStrategy;
 import java.sql.ResultSet;
 import java.sql.SQLException;
 import org.junit.jupiter.api.Test;
 
 /** Integration tests for SQL statement execution. */
-public class ExecutionIntegrationTests extends BaseFakeServiceIntegrationTests {
+public class ExecutionIntegrationTests extends AbstractFakeServiceIntegrationTests {
 
   @Test
   void testInsertStatement() throws SQLException {

src/test/java/com/databricks/jdbc/integration/fakeservice/tests/MetadataIntegrationTests.java

@@ -7,15 +7,15 @@
 import static com.github.tomakehurst.wiremock.client.WireMock.urlEqualTo;
 import static org.junit.jupiter.api.Assertions.*;
 
-import com.databricks.jdbc.integration.fakeservice.BaseFakeServiceIntegrationTests;
+import com.databricks.jdbc.integration.fakeservice.AbstractFakeServiceIntegrationTests;
 import com.github.tomakehurst.wiremock.client.CountMatchingStrategy;
 import java.sql.*;
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 
 /** Integration tests for metadata retrieval. */
-public class MetadataIntegrationTests extends BaseFakeServiceIntegrationTests {
+public class MetadataIntegrationTests extends AbstractFakeServiceIntegrationTests {
 
   private Connection connection;
 

src/test/java/com/databricks/jdbc/integration/fakeservice/tests/MultiChunkExecutionIntegrationTests.java

@@ -4,57 +4,32 @@
 import static com.databricks.jdbc.integration.IntegrationTestUtil.getValidJDBCConnection;
 import static com.github.tomakehurst.wiremock.client.WireMock.getRequestedFor;
 import static com.github.tomakehurst.wiremock.client.WireMock.urlPathMatching;
-import static com.github.tomakehurst.wiremock.core.WireMockConfiguration.wireMockConfig;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 
 import com.databricks.jdbc.core.DatabricksResultSet;
 import com.databricks.jdbc.core.DatabricksResultSetMetaData;
-import com.databricks.jdbc.driver.DatabricksJdbcConstants;
-import com.databricks.jdbc.integration.fakeservice.DatabricksWireMockExtension;
-import com.databricks.jdbc.integration.fakeservice.FakeServiceExtension;
-import com.databricks.jdbc.integration.fakeservice.StubMappingCredentialsCleaner;
-import com.github.tomakehurst.wiremock.extension.Extension;
+import com.databricks.jdbc.integration.fakeservice.AbstractFakeServiceIntegrationTests;
 import java.sql.Connection;
 import java.sql.ResultSet;
 import java.sql.SQLException;
 import java.sql.Statement;
 import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeAll;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
-import org.junit.jupiter.api.extension.RegisterExtension;
 
 /** Test SQL execution with results spanning multiple chunks. */
-public class MultiChunkExecutionIntegrationTests {
-
-  /**
-   * {@link FakeServiceExtension} for {@link DatabricksJdbcConstants.FakeServiceType#SQL_EXEC}.
-   * Intercepts all requests to SQL Execution API.
-   */
-  @RegisterExtension
-  private static final FakeServiceExtension sqlExecApiExtension =
-      new FakeServiceExtension(
-          new DatabricksWireMockExtension.Builder()
-              .options(
-                  wireMockConfig().dynamicPort().dynamicHttpsPort().extensions(getExtensions())),
-          DatabricksJdbcConstants.FakeServiceType.SQL_EXEC,
-          "https://e2-dogfood.staging.cloud.databricks.com");
-
-  /**
-   * {@link FakeServiceExtension} for {@link DatabricksJdbcConstants.FakeServiceType#CLOUD_FETCH}.
-   * Intercepts all requests to Cloud Fetch API.
-   */
-  @RegisterExtension
-  private static final FakeServiceExtension cloudFetchApiExtension =
-      new FakeServiceExtension(
-          new DatabricksWireMockExtension.Builder()
-              .options(
-                  wireMockConfig().dynamicPort().dynamicHttpsPort().extensions(getExtensions())),
-          DatabricksJdbcConstants.FakeServiceType.CLOUD_FETCH,
-          "https://e2-dogfood-core.s3.us-west-2.amazonaws.com");
+public class MultiChunkExecutionIntegrationTests extends AbstractFakeServiceIntegrationTests {
 
   private Connection connection;
 
+  @BeforeAll
+  static void beforeAll() {
+    setSqlExecApiTargetUrl("https://e2-dogfood.staging.cloud.databricks.com");
+    setCloudFetchApiTargetUrl("https://e2-dogfood-core.s3.us-west-2.amazonaws.com");
+  }
+
   @BeforeEach
   void setUp() throws SQLException {
     connection = getValidJDBCConnection();
@@ -94,7 +69,7 @@ void testMultiChunkSelect() throws SQLException {
 
       // The number of cloud fetch calls should be equal to the number of chunks
       final int cloudFetchCalls =
-          cloudFetchApiExtension
+          getCloudFetchApiExtension()
               .countRequestsMatching(getRequestedFor(urlPathMatching(".*")).build())
               .getCount();
       assertEquals(metaData.getChunkCount(), cloudFetchCalls);
@@ -103,14 +78,10 @@ void testMultiChunkSelect() throws SQLException {
       // chunks as first chunk link is already fetched
       final String statementId = ((DatabricksResultSet) rs).statementId();
       final String resultChunkPathRegex = String.format(RESULT_CHUNK_PATH, statementId, ".*");
-      sqlExecApiExtension.verify(
-          (int) (metaData.getChunkCount() - 1),
-          getRequestedFor(urlPathMatching(resultChunkPathRegex)));
+      getSqlExecApiExtension()
+          .verify(
+              (int) (metaData.getChunkCount() - 1),
+              getRequestedFor(urlPathMatching(resultChunkPathRegex)));
     }
   }
-
-  /** Returns the extensions to be used for stubbing. */
-  private static Extension[] getExtensions() {
-    return new Extension[] {new StubMappingCredentialsCleaner()};
-  }
 }

src/test/java/com/databricks/jdbc/integration/fakeservice/tests/PreparedStatementIntegrationTests.java

@@ -4,7 +4,7 @@
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 
-import com.databricks.jdbc.integration.fakeservice.BaseFakeServiceIntegrationTests;
+import com.databricks.jdbc.integration.fakeservice.AbstractFakeServiceIntegrationTests;
 import java.sql.Connection;
 import java.sql.PreparedStatement;
 import java.sql.ResultSet;
@@ -14,7 +14,7 @@
 import org.junit.jupiter.api.Test;
 
 /** Integration tests for PreparedStatement operations. */
-public class PreparedStatementIntegrationTests extends BaseFakeServiceIntegrationTests {
+public class PreparedStatementIntegrationTests extends AbstractFakeServiceIntegrationTests {
 
   private Connection connection;
 

src/test/java/com/databricks/jdbc/integration/fakeservice/tests/ResultSetIntegrationTests.java

@@ -3,7 +3,7 @@
 import static com.databricks.jdbc.integration.IntegrationTestUtil.*;
 import static org.junit.jupiter.api.Assertions.*;
 
-import com.databricks.jdbc.integration.fakeservice.BaseFakeServiceIntegrationTests;
+import com.databricks.jdbc.integration.fakeservice.AbstractFakeServiceIntegrationTests;
 import java.math.BigDecimal;
 import java.sql.Date;
 import java.sql.ResultSet;
@@ -12,7 +12,7 @@
 import org.junit.jupiter.api.Test;
 
 /** Integration tests for ResultSet operations. */
-public class ResultSetIntegrationTests extends BaseFakeServiceIntegrationTests {
+public class ResultSetIntegrationTests extends AbstractFakeServiceIntegrationTests {
 
   @Test
   void testRetrievalOfBasicDataTypes() throws SQLException {