Fix test cases for PENDING health state

Author: andythsu

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

@@ -13,16 +13,17 @@
  */
 package io.trino.gateway.ha;
 
+import io.airlift.json.JsonCodec;
 import io.airlift.log.Logger;
+import io.trino.gateway.ha.clustermonitor.ClusterStats;
+import io.trino.gateway.ha.clustermonitor.TrinoHealthStateType;
 import okhttp3.MediaType;
 import okhttp3.OkHttpClient;
 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.jdbi.v3.core.Handle;
 import org.jdbi.v3.core.Jdbi;
 
@@ -32,7 +33,6 @@
 import java.io.InputStream;
 import java.net.URL;
 import java.nio.file.Paths;
-import java.util.Map;
 import java.util.Random;
 import java.util.Scanner;
 
@@ -71,23 +71,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 +142,31 @@ public static void setUpBackend(
                         .build();
         Response response = httpClient.newCall(request).execute();
         assertThat(response.isSuccessful()).isTrue();
+        TrinoHealthStateType newClusterHealthState = TrinoHealthStateType.PENDING;
+        long startTime = System.currentTimeMillis();
+        long lastExecutionTime = System.currentTimeMillis();
+        // pull cluster health states for 10 seconds
+        // It should be enough as the healthcheck is run every second
+        int timeout = 10 * 1000;
+        while (newClusterHealthState != TrinoHealthStateType.HEALTHY && (lastExecutionTime - startTime) < timeout) {
+            // check the state of newly added cluster every second
+            if (System.currentTimeMillis() - lastExecutionTime <= 1000) {
+                continue;
+            }
+            lastExecutionTime = System.currentTimeMillis();
+            request = new Request.Builder()
+                    .url(String.format("http://localhost:%s/api/public/backends/%s/state", routerPort, name))
+                    .get()
+                    .build();
+            response = httpClient.newCall(request).execute();
+            if (response.isSuccessful()) {
+                JsonCodec<ClusterStats> responseCodec = JsonCodec.jsonCodec(ClusterStats.class);
+                ClusterStats clusterStats = responseCodec.fromJson(response.body().string());
+                newClusterHealthState = clusterStats.healthState();
+                log.info("health state for trino cluster %s is %s", name, newClusterHealthState);
+            }
+        }
+        assertThat(newClusterHealthState).isEqualTo(TrinoHealthStateType.HEALTHY);
     }
 
     public record TestConfig(String configFilePath, String h2DbFilePath)

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

@@ -26,7 +26,11 @@
 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.jetbrains.annotations.NotNull;
 import org.junit.jupiter.api.AfterAll;
 import org.junit.jupiter.api.BeforeAll;
 import org.junit.jupiter.api.Test;
@@ -37,9 +41,12 @@
 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.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;
 
@@ -83,11 +90,26 @@ public void setup()
         int backend2Port = scheduledTrino.getMappedPort(8080);
 
         HaGatewayTestUtils.prepareMockBackend(customBackend, customBackendPort, "default custom response");
-        HaGatewayTestUtils.setPathSpecificResponses(customBackend, ImmutableMap.of(
+        Map<String, String> pathResponseMap = ImmutableMap.of(
                 oauthInitiatePath, oauthInitialResponse,
                 oauthCallbackPath, oauthCallbackResponse,
                 CUSTOM_PATH, CUSTOM_RESPONSE,
-                CUSTOM_LOGOUT, ""));
+                CUSTOM_LOGOUT, "");
+        customBackend.setDispatcher(new Dispatcher() {
+            @Override
+            public MockResponse dispatch(RecordedRequest request)
+            {
+                if (pathResponseMap.containsKey(request.getPath())) {
+                    return new MockResponse().setResponseCode(200).setBody(pathResponseMap.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

@@ -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-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.*'