feat(tests): Refactor test environment to use a single, shared MockServer instance

Author: adutra

oauth2/core/src/test/java/com/dremio/iceberg/authmgr/oauth2/endpoint/EndpointProviderTest.java

@@ -22,12 +22,11 @@
 
 import com.dremio.iceberg.authmgr.oauth2.http.HttpClient;
 import com.dremio.iceberg.authmgr.oauth2.test.TestEnvironment;
-import com.dremio.iceberg.authmgr.oauth2.test.server.UnitTestHttpServer;
+import com.dremio.iceberg.authmgr.oauth2.test.TestServer;
 import com.nimbusds.oauth2.sdk.ParseException;
 import org.junit.jupiter.api.Test;
-import org.junit.jupiter.params.ParameterizedTest;
-import org.junit.jupiter.params.provider.CsvSource;
-import org.junit.jupiter.params.provider.ValueSource;
+import org.junitpioneer.jupiter.cartesian.CartesianTest;
+import org.junitpioneer.jupiter.cartesian.CartesianTest.Values;
 import org.mockserver.model.HttpRequest;
 import org.mockserver.model.HttpResponse;
 import org.mockserver.model.JsonBody;
@@ -50,12 +49,19 @@ void withoutDiscovery() {
     }
   }
 
-  @ParameterizedTest
-  @ValueSource(booleans = {true, false})
-  void withDiscovery(boolean includeDeviceAuthEndpoint) {
+  @CartesianTest
+  void withDiscovery(
+      @Values(booleans = {true, false}) boolean includeDeviceAuthEndpoint,
+      @Values(
+              strings = {
+                ".well-known/openid-configuration",
+                ".well-known/oauth-authorization-server"
+              })
+          String wellKnownPath) {
     try (TestEnvironment env =
         TestEnvironment.builder()
             .includeDeviceAuthEndpointInDiscoveryMetadata(includeDeviceAuthEndpoint)
+            .wellKnownPath(wellKnownPath)
             .build()) {
       EndpointProvider endpointProvider =
           EndpointProvider.create(env.getOAuth2Config(), HttpClient.DEFAULT);
@@ -74,33 +80,6 @@ void withDiscovery(boolean includeDeviceAuthEndpoint) {
     }
   }
 
-  @ParameterizedTest
-  @CsvSource({
-    "''              , /.well-known/openid-configuration",
-    "/               , /.well-known/openid-configuration",
-    "''              , /.well-known/oauth-authorization-server",
-    "/               , /.well-known/oauth-authorization-server",
-    "/realms/master  , /realms/master/.well-known/openid-configuration",
-    "/realms/master/ , /realms/master/.well-known/openid-configuration",
-    "/realms/master  , /realms/master/.well-known/oauth-authorization-server",
-    "/realms/master/ , /realms/master/.well-known/oauth-authorization-server"
-  })
-  void fetchOpenIdProviderMetadataSuccess(String contextPath, String wellKnownPath) {
-    try (TestEnvironment env =
-        TestEnvironment.builder()
-            .authorizationServerContextPath(contextPath)
-            .wellKnownPath(wellKnownPath)
-            .build()) {
-      EndpointProvider endpointProvider =
-          EndpointProvider.create(env.getOAuth2Config(), HttpClient.DEFAULT);
-      var actual = endpointProvider.getOpenIdProviderMetadata();
-      assertThat(actual.getTokenEndpointURI()).isEqualTo(env.getTokenEndpoint());
-      assertThat(actual.getAuthorizationEndpointURI()).isEqualTo(env.getAuthorizationEndpoint());
-      assertThat(actual.getDeviceAuthorizationEndpointURI())
-          .isEqualTo(env.getDeviceAuthorizationEndpoint());
-    }
-  }
-
   @Test
   void fetchOpenIdProviderMetadataWrongEndpoint() {
     try (TestEnvironment env = TestEnvironment.builder().createDefaultExpectations(false).build()) {
@@ -128,9 +107,8 @@ void fetchOpenIdProviderMetadataWrongEndpoint() {
   @Test
   void fetchOpenIdProviderMetadataWrongData() {
     try (TestEnvironment env = TestEnvironment.builder().createDefaultExpectations(false).build()) {
-      ((UnitTestHttpServer) env.getServer())
-          .getClientAndServer()
-          .when(HttpRequest.request())
+      TestServer.getInstance()
+          .when(HttpRequest.request().withPath(env.getAuthorizationServerUrl().getPath() + ".*"))
           .respond(
               HttpResponse.response()
                   .withStatusCode(200)

oauth2/core/src/testFixtures/java/com/dremio/iceberg/authmgr/oauth2/test/TestEnvironment.java

@@ -46,9 +46,6 @@
 import com.dremio.iceberg.authmgr.oauth2.test.expectation.ImmutablePasswordExpectation;
 import com.dremio.iceberg.authmgr.oauth2.test.expectation.ImmutableRefreshTokenExpectation;
 import com.dremio.iceberg.authmgr.oauth2.test.expectation.ImmutableTokenExchangeExpectation;
-import com.dremio.iceberg.authmgr.oauth2.test.server.HttpServer;
-import com.dremio.iceberg.authmgr.oauth2.test.server.IntegrationTestHttpServer;
-import com.dremio.iceberg.authmgr.oauth2.test.server.UnitTestHttpServer;
 import com.dremio.iceberg.authmgr.oauth2.test.user.InteractiveUserEmulator;
 import com.dremio.iceberg.authmgr.oauth2.test.user.UserBehavior;
 import com.dremio.iceberg.authmgr.oauth2.test.user.UserEmulator;
@@ -79,6 +76,7 @@
 import java.util.OptionalInt;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
 import javax.net.ssl.SSLContext;
 import org.apache.iceberg.CatalogProperties;
 import org.apache.iceberg.catalog.SessionCatalog.SessionContext;
@@ -94,6 +92,8 @@
 @Value.Immutable(copy = false)
 public abstract class TestEnvironment implements AutoCloseable {
 
+  private static final AtomicInteger ID_COUNTER = new AtomicInteger();
+
   public static Builder builder() {
     return ImmutableTestEnvironment.builder();
   }
@@ -105,6 +105,11 @@ public void validate() {
     }
   }
 
+  @Value.Default
+  public String getId() {
+    return "env" + ID_COUNTER.incrementAndGet();
+  }
+
   @Value.Default
   public GrantType getGrantType() {
     return GrantType.CLIENT_CREDENTIALS;
@@ -140,11 +145,6 @@ public boolean isCreateDefaultExpectations() {
     return isUnitTest();
   }
 
-  @Value.Lazy
-  public HttpServer getServer() {
-    return isUnitTest() ? new UnitTestHttpServer(isSsl()) : new IntegrationTestHttpServer();
-  }
-
   @Value.Default
   public boolean isSsl() {
     return false;
@@ -161,7 +161,9 @@ public int getExecutorPoolSize() {
   }
 
   public void reset() {
-    getServer().reset();
+    if (isUnitTest()) {
+      TestServer.clear(getId());
+    }
   }
 
   @Override
@@ -175,24 +177,31 @@ public void close() {
     } catch (InterruptedException e) {
       Thread.currentThread().interrupt();
     }
-    getServer().close();
+    reset();
   }
 
   @Value.Default
   public URI getServerRootUrl() {
-    // Note: the default value is for unit tests; integration tests must provide the server root URL
-    // to avoid circular dependencies when creating the TestEnvironment instance
-    return getServer().getRootUrl();
+    if (!isUnitTest()) {
+      throw new IllegalStateException("Server root URL must be provided for integration tests");
+    }
+    return URI.create(
+        (isSsl() ? "https" : "http")
+            + "://localhost:"
+            + TestServer.getInstance().getLocalPort()
+            + "/"
+            + getId()
+            + "/");
   }
 
   @Value.Default
   public String getAuthorizationServerContextPath() {
-    return "/realms/master/";
+    return "realms/master/";
   }
 
   @Value.Default
   public String getCatalogServerContextPath() {
-    return "/api/catalog/";
+    return "api/catalog/";
   }
 
   @Value.Default

…uth2/test/server/UnitTestHttpServer.java …berg/authmgr/oauth2/test/TestServer.javaoauth2/core/src/testFixtures/java/com/dremio/iceberg/authmgr/oauth2/test/server/UnitTestHttpServer.java renamed to oauth2/core/src/testFixtures/java/com/dremio/iceberg/authmgr/oauth2/test/TestServer.java oauth2/core/src/testFixtures/java/com/dremio/iceberg/authmgr/oauth2/test/server/UnitTestHttpServer.java renamed to oauth2/core/src/testFixtures/java/com/dremio/iceberg/authmgr/oauth2/test/TestServer.java

@@ -13,30 +13,25 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package com.dremio.iceberg.authmgr.oauth2.test.server;
+package com.dremio.iceberg.authmgr.oauth2.test;
 
 import java.io.IOException;
-import java.net.URI;
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.nio.file.Paths;
-import java.util.concurrent.atomic.AtomicInteger;
 import org.mockserver.configuration.Configuration;
 import org.mockserver.integration.ClientAndServer;
+import org.mockserver.model.HttpRequest;
 
-public class UnitTestHttpServer implements HttpServer {
+public final class TestServer {
 
-  private static final AtomicInteger COUNTER = new AtomicInteger(1);
+  private static final ClientAndServer INSTANCE;
 
-  private final boolean ssl;
-  private final ClientAndServer clientAndServer;
-
-  public UnitTestHttpServer(boolean ssl) {
-    this.ssl = ssl;
+  static {
     Configuration configuration = Configuration.configuration();
     String outputDir = System.getProperty("authmgr.test.mockserver.memoryUsageCsvDirectory");
     if (outputDir != null) {
-      Path outputPath = Paths.get(outputDir).resolve("server-" + COUNTER.getAndIncrement());
+      Path outputPath = Paths.get(outputDir);
       try {
         Files.createDirectories(outputPath);
       } catch (IOException e) {
@@ -45,25 +40,18 @@ public UnitTestHttpServer(boolean ssl) {
       configuration.outputMemoryUsageCsv(true);
       configuration.memoryUsageCsvDirectory(outputPath.toString());
     }
-    clientAndServer = ClientAndServer.startClientAndServer(configuration);
-  }
-
-  public ClientAndServer getClientAndServer() {
-    return clientAndServer;
+    INSTANCE = ClientAndServer.startClientAndServer(configuration);
+    Runtime.getRuntime().addShutdownHook(new Thread(INSTANCE::close));
   }
 
-  @Override
-  public URI getRootUrl() {
-    return URI.create((ssl ? "https" : "http") + "://localhost:" + clientAndServer.getLocalPort());
-  }
+  private TestServer() {}
 
-  @Override
-  public void reset() {
-    clientAndServer.reset();
+  public static ClientAndServer getInstance() {
+    return INSTANCE;
   }
 
-  @Override
-  public void close() {
-    clientAndServer.close();
+  /** Clears all expectations and responses for the given test environment id. */
+  public static void clear(String testEnvironmentId) {
+    INSTANCE.clear(HttpRequest.request().withPath("/" + testEnvironmentId + "/.*"));
   }
 }

oauth2/core/src/testFixtures/java/com/dremio/iceberg/authmgr/oauth2/test/expectation/AbstractExpectation.java

@@ -16,18 +16,12 @@
 package com.dremio.iceberg.authmgr.oauth2.test.expectation;
 
 import com.dremio.iceberg.authmgr.oauth2.test.TestEnvironment;
-import com.dremio.iceberg.authmgr.oauth2.test.server.UnitTestHttpServer;
 import org.immutables.value.Value;
-import org.mockserver.integration.ClientAndServer;
 
 public abstract class AbstractExpectation {
 
   @Value.Parameter(order = 1)
   protected abstract TestEnvironment getTestEnvironment();
 
   public abstract void create();
-
-  protected ClientAndServer getClientAndServer() {
-    return ((UnitTestHttpServer) getTestEnvironment().getServer()).getClientAndServer();
-  }
 }

oauth2/core/src/testFixtures/java/com/dremio/iceberg/authmgr/oauth2/test/expectation/AuthorizationCodeExpectation.java

@@ -21,6 +21,7 @@
 import static com.dremio.iceberg.authmgr.oauth2.test.TestConstants.SCOPE2;
 import static com.dremio.iceberg.authmgr.oauth2.test.expectation.ErrorExpectation.AUTHORIZATION_SERVER_ERROR_RESPONSE;
 
+import com.dremio.iceberg.authmgr.oauth2.test.TestServer;
 import com.dremio.iceberg.authmgr.tools.immutables.AuthManagerImmutable;
 import com.google.common.collect.ImmutableMap;
 import com.nimbusds.oauth2.sdk.AuthorizationCode;
@@ -120,7 +121,7 @@ private void createAuthEndpointExpectation() {
       request.withQueryStringParameter(
           "code_challenge_method", getTestEnvironment().getCodeChallengeMethod().getValue());
     }
-    getClientAndServer()
+    TestServer.getInstance()
         .when(request)
         .respond(
             httpRequest -> {

oauth2/core/src/testFixtures/java/com/dremio/iceberg/authmgr/oauth2/test/expectation/ConfigEndpointExpectation.java

@@ -17,6 +17,7 @@
 
 import static com.dremio.iceberg.authmgr.oauth2.test.expectation.ExpectationUtils.getJsonBody;
 
+import com.dremio.iceberg.authmgr.oauth2.test.TestServer;
 import com.dremio.iceberg.authmgr.tools.immutables.AuthManagerImmutable;
 import com.google.common.collect.ImmutableList;
 import java.util.HashMap;
@@ -61,7 +62,7 @@ public void create() {
             .withOverrides(new HashMap<>())
             .withEndpoints(endpoints)
             .build();
-    getClientAndServer()
+    TestServer.getInstance()
         .when(
             HttpRequest.request()
                 .withMethod("GET")

oauth2/core/src/testFixtures/java/com/dremio/iceberg/authmgr/oauth2/test/expectation/DeviceCodeExpectation.java

@@ -20,6 +20,7 @@
 import static com.dremio.iceberg.authmgr.oauth2.test.expectation.ErrorExpectation.AUTHORIZATION_SERVER_ERROR_RESPONSE;
 import static org.mockserver.model.Parameter.param;
 
+import com.dremio.iceberg.authmgr.oauth2.test.TestServer;
 import com.dremio.iceberg.authmgr.tools.immutables.AuthManagerImmutable;
 import com.google.common.collect.ImmutableMap;
 import com.nimbusds.oauth2.sdk.GrantType;
@@ -82,7 +83,7 @@ protected HttpResponse response(
   }
 
   private void createDeviceAuthEndpointExpectation() {
-    getClientAndServer()
+    TestServer.getInstance()
         .when(
             HttpRequest.request()
                 .withMethod("POST")
@@ -119,7 +120,7 @@ private void createDeviceAuthEndpointExpectation() {
   private void createDeviceVerificationEndpointExpectation() {
     String path = getTestEnvironment().getDeviceVerificationEndpoint().getPath();
     // Expect the device verification page to be opened in a browser
-    getClientAndServer()
+    TestServer.getInstance()
         .when(HttpRequest.request().withMethod("GET").withPath(path))
         .respond(
             HttpResponse.response()
@@ -137,7 +138,7 @@ private void createDeviceVerificationEndpointExpectation() {
                             + "</body></html>",
                         MediaType.TEXT_HTML)));
     // Expect the device verification code to be sent by the user after opening the page
-    getClientAndServer()
+    TestServer.getInstance()
         .when(
             HttpRequest.request()
                 .withMethod("POST")

oauth2/core/src/testFixtures/java/com/dremio/iceberg/authmgr/oauth2/test/expectation/ErrorExpectation.java

@@ -15,6 +15,7 @@
  */
 package com.dremio.iceberg.authmgr.oauth2.test.expectation;
 
+import com.dremio.iceberg.authmgr.oauth2.test.TestServer;
 import com.dremio.iceberg.authmgr.tools.immutables.AuthManagerImmutable;
 import org.mockserver.model.HttpRequest;
 import org.mockserver.model.HttpResponse;
@@ -42,15 +43,15 @@ public abstract class ErrorExpectation extends AbstractExpectation {
 
   @Override
   public void create() {
-    getClientAndServer()
+    TestServer.getInstance()
         .when(
             HttpRequest.request()
-                .withPath(getTestEnvironment().getAuthorizationServerContextPath() + ".*"))
+                .withPath(getTestEnvironment().getAuthorizationServerUrl().getPath() + ".*"))
         .respond(AUTHORIZATION_SERVER_ERROR_RESPONSE);
-    getClientAndServer()
+    TestServer.getInstance()
         .when(
             HttpRequest.request()
-                .withPath(getTestEnvironment().getCatalogServerContextPath() + ".*"))
+                .withPath(getTestEnvironment().getCatalogServerUrl().getPath() + ".*"))
         .respond(CATALOG_SERVER_ERROR_RESPONSE);
   }
 }

oauth2/core/src/testFixtures/java/com/dremio/iceberg/authmgr/oauth2/test/expectation/InitialTokenFetchExpectation.java

@@ -15,11 +15,13 @@
  */
 package com.dremio.iceberg.authmgr.oauth2.test.expectation;
 
+import com.dremio.iceberg.authmgr.oauth2.test.TestServer;
+
 public abstract class InitialTokenFetchExpectation extends AbstractTokenEndpointExpectation {
 
   @Override
   public void create() {
-    getClientAndServer()
+    TestServer.getInstance()
         .when(request())
         .respond(httpRequest -> response(httpRequest, "access_initial", "refresh_initial"));
   }