Fix test cases for PENDING health state

Author: andythsu

gateway-ha/src/test/java/io/trino/gateway/ha/HaGatewayTestUtils.java

@@ -13,16 +13,19 @@
  */
 package io.trino.gateway.ha;
 
+import com.google.common.base.Stopwatch;
+import io.airlift.json.JsonCodec;
 import io.airlift.log.Logger;
+import io.trino.gateway.ha.clustermonitor.ClusterStats;
+import io.trino.gateway.ha.clustermonitor.TrinoStatus;
 import okhttp3.MediaType;
 import okhttp3.OkHttpClient;
 import okhttp3.Request;
 import okhttp3.RequestBody;
 import okhttp3.Response;
-import okhttp3.mockwebserver.Dispatcher;
+import okhttp3.ResponseBody;
 import okhttp3.mockwebserver.MockResponse;
 import okhttp3.mockwebserver.MockWebServer;
-import okhttp3.mockwebserver.RecordedRequest;
 import org.jdbi.v3.core.Handle;
 import org.jdbi.v3.core.Jdbi;
 
@@ -32,12 +35,13 @@
 import java.io.InputStream;
 import java.net.URL;
 import java.nio.file.Paths;
-import java.util.Map;
+import java.time.Duration;
 import java.util.Random;
 import java.util.Scanner;
 
 import static com.google.common.net.HttpHeaders.CONTENT_ENCODING;
 import static com.google.common.net.MediaType.PLAIN_TEXT_UTF_8;
+import static java.lang.String.format;
 import static java.nio.charset.StandardCharsets.UTF_8;
 import static java.util.Objects.requireNonNull;
 import static org.assertj.core.api.Assertions.assertThat;
@@ -71,23 +75,6 @@ public static void prepareMockBackend(
                 .setResponseCode(200));
     }
 
-    public static void setPathSpecificResponses(
-            MockWebServer backend, Map<String, String> pathResponseMap)
-    {
-        Dispatcher dispatcher = new Dispatcher()
-        {
-            @Override
-            public MockResponse dispatch(RecordedRequest request)
-            {
-                if (pathResponseMap.containsKey(request.getPath())) {
-                    return new MockResponse().setResponseCode(200).setBody(pathResponseMap.get(request.getPath()));
-                }
-                return new MockResponse().setResponseCode(404);
-            }
-        };
-        backend.setDispatcher(dispatcher);
-    }
-
     public static TestConfig buildGatewayConfigAndSeedDb(int routerPort, String configFile)
             throws Exception
     {
@@ -159,6 +146,31 @@ public static void setUpBackend(
                         .build();
         Response response = httpClient.newCall(request).execute();
         assertThat(response.isSuccessful()).isTrue();
+        TrinoStatus newClusterHealthState = TrinoStatus.PENDING;
+        Stopwatch stopwatch = Stopwatch.createStarted();
+        // pull cluster health states for 10 seconds
+        // It should be enough as the healthcheck is run every second
+        Duration timeout = Duration.ofSeconds(10);
+        while (newClusterHealthState != TrinoStatus.HEALTHY && stopwatch.elapsed().compareTo(timeout) < 0) {
+            // check the state of newly added cluster every second
+            Request getBackendStateRequest = new Request.Builder()
+                    .url(format("http://localhost:%s/api/public/backends/%s/state", routerPort, name))
+                    .get()
+                    .build();
+            try (Response getBackendStateResponse = httpClient.newCall(getBackendStateRequest).execute()) {
+                if (getBackendStateResponse.isSuccessful()) {
+                    JsonCodec<ClusterStats> responseCodec = JsonCodec.jsonCodec(ClusterStats.class);
+                    ResponseBody getBackendStateResponseBody = getBackendStateResponse.body();
+                    if (getBackendStateResponseBody != null) {
+                        ClusterStats clusterStats = responseCodec.fromJson(getBackendStateResponseBody.string());
+                        newClusterHealthState = clusterStats.healthState();
+                        log.debug("health state for cluster %s is %s", name, newClusterHealthState);
+                    }
+                }
+            }
+            Thread.sleep(1000);
+        }
+        assertThat(newClusterHealthState).isEqualTo(TrinoStatus.HEALTHY);
     }
 
     public record TestConfig(String configFilePath, String h2DbFilePath)

gateway-ha/src/test/java/io/trino/gateway/ha/TestGatewayHaMultipleBackend.java

@@ -26,7 +26,10 @@
 import okhttp3.Request;
 import okhttp3.RequestBody;
 import okhttp3.Response;
+import okhttp3.mockwebserver.Dispatcher;
+import okhttp3.mockwebserver.MockResponse;
 import okhttp3.mockwebserver.MockWebServer;
+import okhttp3.mockwebserver.RecordedRequest;
 import org.junit.jupiter.api.AfterAll;
 import org.junit.jupiter.api.BeforeAll;
 import org.junit.jupiter.api.Test;
@@ -37,10 +40,13 @@
 import java.io.IOException;
 import java.util.Base64;
 import java.util.List;
+import java.util.Map;
 import java.util.Optional;
 import java.util.concurrent.TimeUnit;
 
 import static com.google.common.collect.MoreCollectors.onlyElement;
+import static com.google.common.net.HttpHeaders.CONTENT_TYPE;
+import static com.google.common.net.MediaType.JSON_UTF_8;
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.testcontainers.utility.MountableFile.forClasspathResource;
 
@@ -84,11 +90,26 @@ void setup()
         int backend2Port = scheduledTrino.getMappedPort(8080);
 
         HaGatewayTestUtils.prepareMockBackend(customBackend, customBackendPort, "default custom response");
-        HaGatewayTestUtils.setPathSpecificResponses(customBackend, ImmutableMap.of(
-                oauthInitiatePath, oauthInitialResponse,
-                oauthCallbackPath, oauthCallbackResponse,
-                CUSTOM_PATH, CUSTOM_RESPONSE,
-                CUSTOM_LOGOUT, ""));
+        customBackend.setDispatcher(new Dispatcher() {
+            @Override
+            public MockResponse dispatch(RecordedRequest request)
+            {
+                Map<String, String> pathResponse = ImmutableMap.of(
+                        oauthInitiatePath, oauthInitialResponse,
+                        oauthCallbackPath, oauthCallbackResponse,
+                        CUSTOM_PATH, CUSTOM_RESPONSE,
+                        CUSTOM_LOGOUT, "");
+                if (pathResponse.containsKey(request.getPath())) {
+                    return new MockResponse().setResponseCode(200).setBody(pathResponse.get(request.getPath()));
+                }
+                if (request.getPath().equals("/v1/info")) {
+                    return new MockResponse().setResponseCode(200)
+                            .setHeader(CONTENT_TYPE, JSON_UTF_8)
+                            .setBody("{\"starting\": false}");
+                }
+                return new MockResponse().setResponseCode(404);
+            }
+        });
 
         // seed database
         HaGatewayTestUtils.TestConfig testConfig =

gateway-ha/src/test/resources/test-config-template.yml

@@ -11,6 +11,17 @@ dataStore:
 
 modules:
   - io.trino.gateway.ha.module.HaGatewayProviderModule
+  - io.trino.gateway.ha.module.ClusterStateListenerModule
+  - io.trino.gateway.ha.module.ClusterStatsMonitorModule
+
+managedApps:
+  - io.trino.gateway.ha.clustermonitor.ActiveClusterMonitor
+
+clusterStatsConfiguration:
+  monitorType: INFO_API
+
+monitor:
+  taskDelaySeconds: 1
 
 extraWhitelistPaths:
   - '/v1/custom.*'

gateway-ha/src/test/resources/test-config-with-routing-template.yml

@@ -10,6 +10,17 @@ dataStore:
 
 modules:
   - io.trino.gateway.ha.module.HaGatewayProviderModule
+  - io.trino.gateway.ha.module.ClusterStateListenerModule
+  - io.trino.gateway.ha.module.ClusterStatsMonitorModule
+
+managedApps:
+  - io.trino.gateway.ha.clustermonitor.ActiveClusterMonitor
+
+clusterStatsConfiguration:
+  monitorType: INFO_API
+
+monitor:
+  taskDelaySeconds: 1
 
 extraWhitelistPaths:
   - '/v1/custom.*'

gateway-ha/src/test/resources/test-config-without-x-forwarded-template.yml

@@ -10,6 +10,17 @@ dataStore:
 
 modules:
   - io.trino.gateway.ha.module.HaGatewayProviderModule
+  - io.trino.gateway.ha.module.ClusterStateListenerModule
+  - io.trino.gateway.ha.module.ClusterStatsMonitorModule
+
+managedApps:
+  - io.trino.gateway.ha.clustermonitor.ActiveClusterMonitor
+
+clusterStatsConfiguration:
+  monitorType: INFO_API
+
+monitor:
+  taskDelaySeconds: 1
 
 extraWhitelistPaths:
   - '/v1/custom.*'