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 58e68c2a5..9a3a8cb73 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 @@ -163,7 +163,7 @@ public void testQueryDeliveryToMultipleRoutingGroups() public void testTrinoClusterHostCookie() throws Exception { - RequestBody requestBody = RequestBody.create("SELECT 1", MediaType.get("application/json; charset=utf-8")); + RequestBody requestBody = RequestBody.create("SELECT 1", MEDIA_TYPE); // When X-Trino-Routing-Group is set in header, query should be routed to cluster under the routing group Request requestWithoutCookie = @@ -173,11 +173,12 @@ public void testTrinoClusterHostCookie() .post(requestBody) .addHeader("X-Trino-Routing-Group", "scheduled") .build(); - Response responseOne = httpClient.newCall(requestWithoutCookie).execute(); - assertThat(responseOne.body().string()).contains("http://localhost:" + routerPort); - List cookies = Cookie.parseAll(responseOne.request().url(), responseOne.headers()); - Cookie cookie = cookies.stream().filter(c -> c.name().equals("trinoClusterHost")).collect(onlyElement()); - assertThat(cookie.value()).isEqualTo("localhost"); + Response responseWithoutCookie = httpClient.newCall(requestWithoutCookie).execute(); + assertThat(responseWithoutCookie.body().string()).contains("http://localhost:" + routerPort); + List cookies = Cookie.parseAll(responseWithoutCookie.request().url(), responseWithoutCookie.headers()); + Cookie clusterHostCookie = cookies.stream().filter(c -> c.name().equals("trinoClusterHost")).collect(onlyElement()); + assertThat(clusterHostCookie.value()).isEqualTo("localhost"); + // test with sending the request which includes trinoClusterHost in the cookie // when X-Trino-Routing-Group is set in header, query should be routed to cluster under the routing group Request requestWithCookie = @@ -188,11 +189,11 @@ public void testTrinoClusterHostCookie() .addHeader("X-Trino-Routing-Group", "scheduled") .addHeader("Cookie", "trinoClientHost=foo.example.com") .build(); - Response responseTwo = httpClient.newCall(requestWithCookie).execute(); - assertThat(responseTwo.body().string()).contains("http://localhost:" + routerPort); - cookies = Cookie.parseAll(responseTwo.request().url(), responseTwo.headers()); - cookie = cookies.stream().filter(c -> c.name().equals("trinoClusterHost")).collect(onlyElement()); - assertThat(cookie.value()).isEqualTo("localhost"); + Response responseWithCookie = httpClient.newCall(requestWithCookie).execute(); + assertThat(responseWithCookie.body().string()).contains("http://localhost:" + routerPort); + List overridenCookies = Cookie.parseAll(responseWithCookie.request().url(), responseWithCookie.headers()); + Cookie overridenClusterHostCookie = overridenCookies.stream().filter(c -> c.name().equals("trinoClusterHost")).collect(onlyElement()); + assertThat(overridenClusterHostCookie.value()).isEqualTo("localhost"); } @Test