From b72f8cad3bbdb90ea6cb4d484f6e13bc1a180443 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 | 56 ++++++++++++------- .../ha/TestGatewayHaMultipleBackend.java | 31 ++++++++-- .../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, 95 insertions(+), 25 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 ea390a68d..5c2a679c9 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,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,16 @@ 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 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 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 +78,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 +149,32 @@ public static void setUpBackend( .build(); Response response = httpClient.newCall(request).execute(); assertThat(response.isSuccessful()).isTrue(); + verifyTrinoStatus(routerPort, name); + } + + private static void verifyTrinoStatus(int port, String name) + throws IOException + { + Request getBackendStateRequest = new Request.Builder() + .url(format("http://localhost:%s/api/public/backends/%s/state", port, name)) + .get() + .build(); + + int tries = 10; + for (int i = 0; i < tries; i++) { + // check the state of newly added cluster every second + try (Response getBackendStateResponse = httpClient.newCall(getBackendStateRequest).execute()) { + checkState(getBackendStateResponse.isSuccessful()); + JsonCodec responseCodec = JsonCodec.jsonCodec(ClusterStats.class); + ResponseBody getBackendStateResponseBody = requireNonNull(getBackendStateResponse.body(), "getBackendStateResponse.body() is null"); + ClusterStats clusterStats = responseCodec.fromJson(getBackendStateResponseBody.string()); + if (clusterStats.trinoStatus() == TrinoStatus.HEALTHY) { + return; + } + } + sleepUninterruptibly(1, TimeUnit.SECONDS); + } + throw new IllegalStateException("Trino cluster is not 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 af680a6a7..56ebfb5eb 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,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 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 = diff --git a/gateway-ha/src/test/resources/test-config-template.yml b/gateway-ha/src/test/resources/test-config-template.yml index beafe30c7..81ea1b461 100644 --- a/gateway-ha/src/test/resources/test-config-template.yml +++ b/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.*' 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.*'