From a2b693a62b411529b7ebca9839df2e8df4620ab4 Mon Sep 17 00:00:00 2001 From: andythsu Date: Wed, 25 Sep 2024 09:44:38 -0400 Subject: [PATCH] Fix test cases for PENDING health state --- .../trino/gateway/ha/HaGatewayTestUtils.java | 48 +++++++++++-------- .../ha/TestGatewayHaMultipleBackend.java | 26 +++++++++- .../test/resources/test-config-template.yml | 11 +++++ .../test-config-with-routing-template.yml | 11 +++++ ...st-config-without-x-forwarded-template.yml | 11 +++++ 5 files changed, 85 insertions(+), 22 deletions(-) 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 index 9a81345b8..48c4f0aaf 100644 --- 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,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 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 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) diff --git a/gateway-ha/src/test/java/io/trino/gateway/ha/TestGatewayHaMultipleBackend.java b/gateway-ha/src/test/java/io/trino/gateway/ha/TestGatewayHaMultipleBackend.java index 1ffcc6dd5..900730d6b 100644 --- a/gateway-ha/src/test/java/io/trino/gateway/ha/TestGatewayHaMultipleBackend.java +++ b/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; @@ -81,11 +88,26 @@ public void setup() int backend2Port = scheduledTrino.getMappedPort(8080); HaGatewayTestUtils.prepareMockBackend(customBackend, customBackendPort, "default custom response"); - HaGatewayTestUtils.setPathSpecificResponses(customBackend, ImmutableMap.of( + Map 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 = diff --git a/gateway-ha/src/test/resources/test-config-template.yml b/gateway-ha/src/test/resources/test-config-template.yml index b22e27fe9..f080fdb4b 100644 --- a/gateway-ha/src/test/resources/test-config-template.yml +++ b/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.*' 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 index 84a5b0f32..817b93e20 100644 --- a/gateway-ha/src/test/resources/test-config-with-routing-template.yml +++ b/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.*' 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 index 15eb4b0cc..4be161476 100644 --- 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 @@ -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.*'