Allow replacing test config with environment variable

andythsu · ebyhr · commit c3a84c2713c6 · 2025-06-21T09:28:46.000+09:00
diff --git a/gateway-ha/src/test/java/io/trino/gateway/ha/HaGatewayTestUtils.java b/gateway-ha/src/test/java/io/trino/gateway/ha/HaGatewayTestUtils.java
@@ -13,6 +13,7 @@
  */
 package io.trino.gateway.ha;
 
+import com.google.common.collect.ImmutableMap;
 import io.airlift.json.JsonCodec;
 import io.airlift.log.Logger;
 import io.trino.gateway.ha.clustermonitor.ClusterStats;
@@ -34,16 +35,17 @@
 import java.io.File;
 import java.io.IOException;
 import java.io.InputStream;
-import java.net.URL;
 import java.nio.file.Files;
 import java.nio.file.Path;
+import java.util.Map;
 import java.util.Scanner;
 import java.util.concurrent.TimeUnit;
 
 import static com.google.common.base.Preconditions.checkState;
 import static com.google.common.net.HttpHeaders.CONTENT_ENCODING;
 import static com.google.common.net.MediaType.PLAIN_TEXT_UTF_8;
 import static com.google.common.util.concurrent.Uninterruptibles.sleepUninterruptibly;
+import static io.trino.gateway.ha.util.ConfigurationUtils.replaceEnvironmentVariables;
 import static java.lang.String.format;
 import static java.nio.charset.StandardCharsets.UTF_8;
 import static java.util.Objects.requireNonNull;
@@ -81,18 +83,23 @@ public static void seedRequiredData(String h2DbFilePath)
     public static File buildGatewayConfig(PostgreSQLContainer postgreSqlContainer, int routerPort, String configFile)
             throws Exception
     {
-        URL resource = HaGatewayTestUtils.class.getClassLoader().getResource("auth/localhost.jks");
-        String configStr =
-                getResourceFileContent(configFile)
-                        .replace("REQUEST_ROUTER_PORT", String.valueOf(routerPort))
-                        .replace("POSTGRESQL_JDBC_URL", postgreSqlContainer.getJdbcUrl())
-                        .replace("POSTGRESQL_USER", postgreSqlContainer.getUsername())
-                        .replace("POSTGRESQL_PASSWORD", postgreSqlContainer.getPassword())
-                        .replace(
-                                "APPLICATION_CONNECTOR_PORT", String.valueOf(30000 + (int) (Math.random() * 1000)))
-                        .replace("ADMIN_CONNECTOR_PORT", String.valueOf(31000 + (int) (Math.random() * 1000)))
-                        .replace("LOCALHOST_JKS", Path.of(resource.toURI()).toString())
-                        .replace("RESOURCES_DIR", Path.of("src", "test", "resources").toAbsolutePath().toString());
+        Map<String, String> additionalVars = ImmutableMap.<String, String>builder()
+                .put("REQUEST_ROUTER_PORT", String.valueOf(routerPort))
+                .putAll(buildPostgresVars(postgreSqlContainer))
+                .buildOrThrow();
+        return buildGatewayConfig(configFile, additionalVars);
+    }
+
+    public static File buildGatewayConfig(String configFile, Map<String, String> additionalVars)
+            throws Exception
+    {
+        Map<String, String> vars = ImmutableMap.<String, String>builder()
+                .put("APPLICATION_CONNECTOR_PORT", String.valueOf(30000 + (int) (Math.random() * 1000)))
+                .put("ADMIN_CONNECTOR_PORT", String.valueOf(31000 + (int) (Math.random() * 1000)))
+                .put("RESOURCES_DIR", Path.of("src", "test", "resources").toAbsolutePath().toString())
+                .putAll(additionalVars)
+                .buildOrThrow();
+        String configStr = replaceEnvironmentVariables(getResourceFileContent(configFile), vars);
 
         File target = File.createTempFile("config-" + System.currentTimeMillis(), "config.yaml");
 
@@ -104,6 +111,15 @@ public static File buildGatewayConfig(PostgreSQLContainer postgreSqlContainer, i
         return target;
     }
 
+    public static Map<String, String> buildPostgresVars(PostgreSQLContainer<?> postgresql)
+    {
+        return ImmutableMap.<String, String>builder()
+                .put("POSTGRESQL_JDBC_URL", postgresql.getJdbcUrl())
+                .put("POSTGRESQL_USER", postgresql.getUsername())
+                .put("POSTGRESQL_PASSWORD", postgresql.getPassword())
+                .buildOrThrow();
+    }
+
     public static String getResourceFileContent(String fileName)
     {
         StringBuilder sb = new StringBuilder();
diff --git a/gateway-ha/src/test/java/io/trino/gateway/ha/security/TestOIDC.java b/gateway-ha/src/test/java/io/trino/gateway/ha/security/TestOIDC.java
@@ -16,6 +16,7 @@
 import com.fasterxml.jackson.core.JsonProcessingException;
 import com.fasterxml.jackson.databind.JsonNode;
 import com.fasterxml.jackson.databind.ObjectMapper;
+import com.google.common.collect.ImmutableMap;
 import io.trino.gateway.ha.HaGatewayLauncher;
 import io.trino.gateway.ha.HaGatewayTestUtils;
 import okhttp3.Cookie;
@@ -44,11 +45,15 @@
 import java.io.File;
 import java.net.CookieManager;
 import java.net.CookiePolicy;
+import java.net.URL;
+import java.nio.file.Path;
 import java.security.SecureRandom;
 import java.security.cert.X509Certificate;
 import java.util.List;
+import java.util.Map;
 import java.util.Optional;
 
+import static io.trino.gateway.ha.HaGatewayTestUtils.buildPostgresVars;
 import static io.trino.gateway.ha.security.OidcCookie.OIDC_COOKIE;
 import static java.lang.String.format;
 import static org.assertj.core.api.Assertions.assertThat;
@@ -144,8 +149,14 @@ void setup()
         PostgreSQLContainer gatewayBackendDatabase = new PostgreSQLContainer("postgres:16");
         gatewayBackendDatabase.start();
 
+        URL resource = HaGatewayTestUtils.class.getClassLoader().getResource("auth/localhost.jks");
+        Map<String, String> additionalVars = ImmutableMap.<String, String>builder()
+                .put("REQUEST_ROUTER_PORT", String.valueOf(ROUTER_PORT))
+                .put("LOCALHOST_JKS", Path.of(resource.toURI()).toString())
+                .putAll(buildPostgresVars(gatewayBackendDatabase))
+                .buildOrThrow();
         File testConfigFile =
-                HaGatewayTestUtils.buildGatewayConfig(gatewayBackendDatabase, ROUTER_PORT, "auth/oauth-test-config.yml");
+                HaGatewayTestUtils.buildGatewayConfig("auth/oauth-test-config.yml", additionalVars);
         String[] args = {testConfigFile.getAbsolutePath()};
         System.out.println(ROUTER_PORT);
         HaGatewayLauncher.main(args);
diff --git a/gateway-ha/src/test/resources/auth/auth-test-config.yml b/gateway-ha/src/test/resources/auth/auth-test-config.yml
@@ -1,11 +1,11 @@
 serverConfig:
   node.environment: test
-  http-server.http.port: REQUEST_ROUTER_PORT
+  http-server.http.port: ${ENV:REQUEST_ROUTER_PORT}
 
 dataStore:
-  jdbcUrl: POSTGRESQL_JDBC_URL
-  user: POSTGRESQL_USER
-  password: POSTGRESQL_PASSWORD
+  jdbcUrl: ${ENV:POSTGRESQL_JDBC_URL}
+  user: ${ENV:POSTGRESQL_USER}
+  password: ${ENV:POSTGRESQL_PASSWORD}
   driver: org.postgresql.Driver
 
 extraWhitelistPaths:
diff --git a/gateway-ha/src/test/resources/auth/oauth-test-config.yml b/gateway-ha/src/test/resources/auth/oauth-test-config.yml
@@ -2,14 +2,14 @@ serverConfig:
   node.environment: test
   http-server.http.port: 8080
   http-server.https.enabled: true
-  http-server.https.port: REQUEST_ROUTER_PORT
-  http-server.https.keystore.path: LOCALHOST_JKS
+  http-server.https.port: ${ENV:REQUEST_ROUTER_PORT}
+  http-server.https.keystore.path: ${ENV:LOCALHOST_JKS}
   http-server.https.keystore.key: 123456
 
 dataStore:
-  jdbcUrl: POSTGRESQL_JDBC_URL
-  user: POSTGRESQL_USER
-  password: POSTGRESQL_PASSWORD
+  jdbcUrl: ${ENV:POSTGRESQL_JDBC_URL}
+  user: ${ENV:POSTGRESQL_USER}
+  password: ${ENV:POSTGRESQL_PASSWORD}
   driver: org.postgresql.Driver
 
 extraWhitelistPaths:
@@ -33,8 +33,8 @@ authentication:
     tokenEndpoint: http://localhost:4444/oauth2/token
     authorizationEndpoint: http://localhost:4444/oauth2/auth
     jwkEndpoint: http://localhost:4444/.well-known/jwks.json
-    redirectUrl: https://localhost:REQUEST_ROUTER_PORT/oidc/callback
-    redirectWebUrl: https://localhost:REQUEST_ROUTER_PORT/
+    redirectUrl: https://localhost:${ENV:REQUEST_ROUTER_PORT}/oidc/callback
+    redirectWebUrl: https://localhost:${ENV:REQUEST_ROUTER_PORT}/
     userIdField: sub
     scopes:
       - openid
diff --git a/gateway-ha/src/test/resources/test-config-template.yml b/gateway-ha/src/test/resources/test-config-template.yml
@@ -1,12 +1,12 @@
 serverConfig:
   node.environment: test
-  http-server.http.port: REQUEST_ROUTER_PORT
+  http-server.http.port: ${ENV:REQUEST_ROUTER_PORT}
 
 includeClusterHostInResponse: true
 dataStore:
-  jdbcUrl: POSTGRESQL_JDBC_URL
-  user: POSTGRESQL_USER
-  password: POSTGRESQL_PASSWORD
+  jdbcUrl: ${ENV:POSTGRESQL_JDBC_URL}
+  user: ${ENV:POSTGRESQL_USER}
+  password: ${ENV:POSTGRESQL_PASSWORD}
   driver: org.postgresql.Driver
 
 clusterStatsConfiguration:
diff --git a/gateway-ha/src/test/resources/test-config-with-routing-rules-api.yml b/gateway-ha/src/test/resources/test-config-with-routing-rules-api.yml
@@ -1,11 +1,11 @@
 serverConfig:
   node.environment: test
-  http-server.http.port: REQUEST_ROUTER_PORT
+  http-server.http.port: ${ENV:REQUEST_ROUTER_PORT}
 
 dataStore:
-  jdbcUrl: POSTGRESQL_JDBC_URL
-  user: POSTGRESQL_USER
-  password: POSTGRESQL_PASSWORD
+  jdbcUrl: ${ENV:POSTGRESQL_JDBC_URL}
+  user: ${ENV:POSTGRESQL_USER}
+  password: ${ENV:POSTGRESQL_PASSWORD}
   driver: org.postgresql.Driver
 
 clusterStatsConfiguration:
@@ -35,4 +35,4 @@ uiConfiguration:
 
 routingRules:
   rulesEngineEnabled: true
-  rulesConfigPath: "RESOURCES_DIR/rules/routing_rules_update.yml"
+  rulesConfigPath: "${ENV:RESOURCES_DIR}/rules/routing_rules_update.yml"
diff --git a/gateway-ha/src/test/resources/test-config-with-routing-template.yml b/gateway-ha/src/test/resources/test-config-with-routing-template.yml
@@ -1,11 +1,11 @@
 serverConfig:
   node.environment: test
-  http-server.http.port: REQUEST_ROUTER_PORT
+  http-server.http.port: ${ENV:REQUEST_ROUTER_PORT}
 
 dataStore:
-  jdbcUrl: POSTGRESQL_JDBC_URL
-  user: POSTGRESQL_USER
-  password: POSTGRESQL_PASSWORD
+  jdbcUrl: ${ENV:POSTGRESQL_JDBC_URL}
+  user: ${ENV:POSTGRESQL_USER}
+  password: ${ENV:POSTGRESQL_PASSWORD}
   driver: org.postgresql.Driver
 
 clusterStatsConfiguration:
@@ -31,4 +31,4 @@ requestAnalyzerConfig:
 
 routingRules:
   rulesEngineEnabled: true
-  rulesConfigPath: "RESOURCES_DIR/rules/routing_rules_trino_query_properties.yml"
+  rulesConfigPath: "${ENV:RESOURCES_DIR}/rules/routing_rules_trino_query_properties.yml"
diff --git a/gateway-ha/src/test/resources/test-config-without-x-forwarded-template.yml b/gateway-ha/src/test/resources/test-config-without-x-forwarded-template.yml
@@ -1,11 +1,11 @@
 serverConfig:
   node.environment: test
-  http-server.http.port: REQUEST_ROUTER_PORT
+  http-server.http.port: ${ENV:REQUEST_ROUTER_PORT}
 
 dataStore:
-  jdbcUrl: POSTGRESQL_JDBC_URL
-  user: POSTGRESQL_USER
-  password: POSTGRESQL_PASSWORD
+  jdbcUrl: ${ENV:POSTGRESQL_JDBC_URL}
+  user: ${ENV:POSTGRESQL_USER}
+  password: ${ENV:POSTGRESQL_PASSWORD}
   driver: org.postgresql.Driver
 
 clusterStatsConfiguration:
